"""Extend ProteinIQ's human-banana OMA count with sequence identity and a species ladder.

Python 3.11+ and Biopython 1.85.

    python sequence_analysis.py           Recompute sequence_summary.json from the saved snapshot.
    python sequence_analysis.py --fetch   Retrieve a fresh sequence snapshot in an empty directory.

The gene-level counts come from analyze.mjs and snapshot.json.gz. This script adds:
1. For every human main-isoform gene with a banana ortholog, the banana ortholog with the
   highest OMA alignment score, and the protein and codon-aligned coding-DNA identity of that pair.
2. A shuffled control: the same human proteins aligned to randomly chosen, non-orthologous
   banana proteins from the same set.
3. The share of human protein-coding sequence and of the human genome in those matched genes.
4. The fraction of human genes with an OMA ortholog in ten other species, counted the same way.
"""

import gzip
import hashlib
import json
import random
import re
import statistics
import sys
import time
import urllib.request
from pathlib import Path

from Bio.Align import PairwiseAligner, substitution_matrices
from Bio.Seq import Seq

DIRECTORY = Path(__file__).resolve().parent
# Split in three so each file stays under the repository's 5 MiB limit.
SNAPSHOT = DIRECTORY / 'sequence_snapshot.json.gz'
SEQUENCE_FILES = {'HUMAN': DIRECTORY / 'sequences_human.json.gz', 'MUSAM': DIRECTORY / 'sequences_banana.json.gz'}
GENE_SNAPSHOT = DIRECTORY / 'snapshot.json.gz'
ORIGIN = 'https://omabrowser.org'
LADDER = ['PANTR', 'MOUSE', 'CHICK', 'DANRE', 'DROME', 'CAEEL', 'ARATH', 'ORYSJ', 'MUSAM', 'YEAST', 'SCHPO']
# T2T-CHM13v2.0 haploid assembly without chrY (Nurk et al. 2022), in base pairs.
HUMAN_GENOME_BP = 3_054_815_472
SEED = 20260925


def get(url, data=None):
    body = None if data is None else json.dumps(data).encode()
    request = urllib.request.Request(url, data=body, headers={'Content-Type': 'application/json', 'User-Agent': 'ProteinIQ-analysis/1.0 (https://www.proteiniq.io)'})
    for attempt in range(5):
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                text = response.read().decode()
                return json.loads(text), dict(response.headers), hashlib.sha256(text.encode()).hexdigest()
        except Exception:
            if attempt == 4:
                raise
            time.sleep(2 ** attempt)


def pages(url, project, log):
    rows, total = [], None
    while url:
        data, headers, digest = get(url)
        count = int(headers['X-Total-Count'])
        total = count if total is None else total
        assert count == total, 'Dataset changed during pagination'
        rows.extend(project(item) for item in data)
        log.append({'url': url, 'sha256': digest, 'total_count': count})
        match = re.search(r'<([^>]+)>;\s*rel="next"', headers.get('Link', ''))
        url = match.group(1) if match else None
        print(f'{len(rows)}/{total}', file=sys.stderr)
    assert len(rows) == total, 'Incomplete pagination'
    return rows


def retrieve():
    assert not SNAPSHOT.exists() and not any(f.exists() for f in SEQUENCE_FILES.values()),         'Preserve the existing snapshot. Fetch in a new directory.'
    log = []
    version = get(f'{ORIGIN}/api/version/')[0]

    human_proteins = pages(
        f'{ORIGIN}/api/genome/HUMAN/proteins/?per_page=10000',
        lambda e: [e['omaid'], e['is_main_isoform'], e['sequence_length'], e['chromosome']],
        log,
    )
    ladder = {}
    for code in LADDER:
        ladder[code] = pages(
            f'{ORIGIN}/api/pairs/HUMAN/{code}/?type=vps&per_page=10000',
            lambda p: [p['entry_1']['omaid'], p['entry_2']['omaid'], p['entry_1']['is_main_isoform'],
                       p['entry_2']['is_main_isoform'], p['rel_type'], p['distance'], p['score']],
            log,
        )

    # Sequences for every main-isoform endpoint of a human-banana pair; the analysis picks the
    # best-scoring partner and draws the shuffled control from the same banana proteins.
    ids = sorted({omaid for h, b, hm, bm, *_ in ladder['MUSAM'] for omaid, main in ((h, hm), (b, bm)) if main})

    sequences = {}
    for start in range(0, len(ids), 100):
        batch = ids[start:start + 100]
        data, _, digest = get(f'{ORIGIN}/api/protein/bulk_retrieve/', {'ids': batch})
        log.append({'url': f'{ORIGIN}/api/protein/bulk_retrieve/', 'ids': [batch[0], batch[-1], len(batch)], 'sha256': digest})
        for item in data:
            target = item['target']
            sequences[target['omaid']] = [target['sequence'], target['cdna']]
        print(f'sequences {len(sequences)}/{len(ids)}', file=sys.stderr)
    assert set(sequences) == set(ids)

    assert get(f'{ORIGIN}/api/version/')[0] == version, 'OMA release changed during retrieval'
    snapshot = {
        'schema_version': 1,
        'version': version,
        'retrieved_at': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
        'human_protein_columns': ['omaid', 'is_main_isoform', 'sequence_length', 'chromosome'],
        'pair_columns': ['human_omaid', 'other_omaid', 'human_is_main_isoform', 'other_is_main_isoform',
                         'rel_type', 'distance', 'score'],
        'sequence_columns': ['protein', 'cdna'],
        'requests': log,
        'human_proteins': human_proteins,
        'pairs': ladder,
    }
    SNAPSHOT.write_bytes(gzip.compress(json.dumps(snapshot).encode(), 9))
    for code, path in SEQUENCE_FILES.items():
        part = {k: v for k, v in sorted(sequences.items()) if k.startswith(code)}
        path.write_bytes(gzip.compress(json.dumps(part).encode(), 9))


def load():
    raw = {'snapshot': SNAPSHOT.read_bytes(), **{p.name: p.read_bytes() for p in SEQUENCE_FILES.values()}}
    snapshot = json.loads(gzip.decompress(raw['snapshot']))
    snapshot['sequences'] = {}
    for path in SEQUENCE_FILES.values():
        snapshot['sequences'].update(json.loads(gzip.decompress(raw[path.name])))
    hashes = {name if name != 'snapshot' else SNAPSHOT.name: hashlib.sha256(data).hexdigest() for name, data in raw.items()}
    return snapshot, hashes


def aligner():
    a = PairwiseAligner()
    a.mode = 'global'
    a.substitution_matrix = substitution_matrices.load('BLOSUM62')
    a.open_gap_score = -11
    a.extend_gap_score = -1
    # Free end gaps, so differences in protein length do not dilute identity.
    a.end_gap_score = 0
    return a


def compare(a, human, banana, sequences):
    hp, hc = sequences[human]
    bp, bc = sequences[banana]
    alignment = a.align(hp, bp)[0]
    columns = identical = 0
    codon_positions = codon_identical = 0
    codons_ok = coding_matches(hp, hc) and coding_matches(bp, bc)
    for (h_start, h_end), (b_start, b_end) in zip(*alignment.aligned):
        for offset in range(h_end - h_start):
            i, j = h_start + offset, b_start + offset
            columns += 1
            identical += hp[i] == bp[j]
            if codons_ok:
                h_codon, b_codon = hc[3 * i:3 * i + 3], bc[3 * j:3 * j + 3]
                codon_positions += 3
                codon_identical += sum(x == y for x, y in zip(h_codon, b_codon))
    return {
        'human': human,
        'banana': banana,
        'human_length': len(hp),
        'aligned_residues': columns,
        'identical_residues': identical,
        'protein_identity': identical / columns if columns else 0,
        'human_coverage': columns / len(hp),
        'codons_ok': codons_ok,
        'aligned_nucleotides': codon_positions,
        'identical_nucleotides': codon_identical,
    }


def coding_matches(protein, cdna):
    """The cDNA must translate to the protein so codons can be placed on the protein alignment."""
    if len(cdna) < 3 * len(protein):
        return False
    translated = str(Seq(cdna[:3 * len(protein)]).translate())
    mismatches = sum(x != y and y != 'X' and x != '*' for x, y in zip(translated, protein))
    return mismatches == 0


def summarize(values):
    values = sorted(values)
    q = statistics.quantiles(values, n=4)
    return {'n': len(values), 'mean': statistics.fmean(values), 'median': statistics.median(values),
            'q1': q[0], 'q3': q[2]}


def analyze(snapshot):
    human = {row[0]: row for row in snapshot['human_proteins']}
    main = {omaid for omaid, is_main, *_ in snapshot['human_proteins'] if is_main}
    assert len(main) == 20430, 'Denominator differs from snapshot.json.gz'
    sequences = snapshot['sequences']

    ladder = {}
    for code, rows in snapshot['pairs'].items():
        matched = {h for h, _, hm, *_ in rows if hm}
        assert matched <= main
        ladder[code] = {'pairs': len(rows), 'human_genes_with_ortholog': len(matched),
                        'percent': 100 * len(matched) / len(main)}

    banana_rows = snapshot['pairs']['MUSAM']
    best = {}
    for h, b, hm, bm, rel_type, _, score in banana_rows:
        if hm and bm and (h not in best or score > best[h][1] or (score == best[h][1] and b < best[h][0])):
            best[h] = (b, score, rel_type)
    assert len(best) == 3683, 'Matched human genes differ from snapshot.json.gz'
    partners = {}
    for h, b, hm, bm, *_ in banana_rows:
        if hm and bm:
            partners.setdefault(h, set()).add(b)

    a = aligner()
    results = [dict(compare(a, h, best[h][0], sequences), rel_type=best[h][2]) for h in sorted(best)]
    one_to_one = [r for r in results if r['rel_type'] == '1:1']

    rng = random.Random(SEED)
    banana_pool = sorted({b for _, b, _, bm, *_ in banana_rows if bm})
    control = []
    for h in sorted(best):
        b = rng.choice(banana_pool)
        while b in partners[h]:
            b = rng.choice(banana_pool)
        control.append(compare(a, h, b, sequences))

    coded = [r for r in results if r['codons_ok'] and r['aligned_nucleotides']]
    control_coded = [r for r in control if r['codons_ok'] and r['aligned_nucleotides']]
    # Coding sequence length = 3 nucleotides per residue plus the stop codon.
    all_cds = sum(3 * (human[h][2] + 1) for h in main)
    matched_cds = sum(3 * (human[h][2] + 1) for h in best)
    aligned_nt = sum(r['aligned_nucleotides'] for r in coded)
    identical_nt = sum(r['identical_nucleotides'] for r in coded)
    coded_cds = sum(3 * (r['human_length'] + 1) for r in coded)
    # Scale the codon-checked subset up to all matched genes by coding length.
    scale = matched_cds / coded_cds

    return {
        'oma_release': snapshot['version']['oma_version'],
        'retrieved_at': snapshot['retrieved_at'],
        'ladder': ladder,
        'banana_pairs': {
            'human_genes': len(best),
            'single_banana_partner': sum(len(partners[h]) == 1 for h in best),
            'best_pair_one_to_one': len(one_to_one),
            'protein_identity': summarize([r['protein_identity'] for r in results]),
            # Ensembl-style: identical residues over the full human protein length.
            'identity_per_human_length': summarize([r['identical_residues'] / r['human_length'] for r in results]),
            'one_to_one_protein_identity': summarize([r['protein_identity'] for r in one_to_one]),
            'one_to_one_identity_per_human_length': summarize([r['identical_residues'] / r['human_length'] for r in one_to_one]),
            'human_coverage': summarize([r['human_coverage'] for r in results]),
            'pooled_protein_identity': sum(r['identical_residues'] for r in results) / sum(r['aligned_residues'] for r in results),
            'codon_checked_pairs': len(coded),
            'nucleotide_identity': summarize([r['identical_nucleotides'] / r['aligned_nucleotides'] for r in coded]),
            'pooled_nucleotide_identity': identical_nt / aligned_nt,
            # Identity over aligned positions in 10-point bins; the last bin includes 100%.
            'protein_identity_histogram': {f'{lo}-{lo + 10}': sum(lo <= min(99.999, 100 * r['protein_identity']) < lo + 10 for r in results) for lo in range(0, 100, 10)},
            'share_above_40_percent_protein_identity': sum(r['protein_identity'] >= 0.4 for r in results) / len(results),
        },
        'shuffled_control': {
            'protein_identity': summarize([r['protein_identity'] for r in control]),
            'identity_per_human_length': summarize([r['identical_residues'] / r['human_length'] for r in control]),
            'nucleotide_identity': summarize([r['identical_nucleotides'] / r['aligned_nucleotides'] for r in control_coded]),
            'human_coverage': summarize([r['human_coverage'] for r in control]),
        },
        'genome_fraction': {
            'human_genome_bp': HUMAN_GENOME_BP,
            'all_main_isoform_cds_bp': all_cds,
            'matched_gene_cds_bp': matched_cds,
            'aligned_nucleotides_scaled_bp': aligned_nt * scale,
            'identical_nucleotides_scaled_bp': identical_nt * scale,
            'percent_genome_all_cds': 100 * all_cds / HUMAN_GENOME_BP,
            'percent_genome_matched_cds': 100 * matched_cds / HUMAN_GENOME_BP,
            'percent_genome_aligned': 100 * aligned_nt * scale / HUMAN_GENOME_BP,
            'percent_genome_identical': 100 * identical_nt * scale / HUMAN_GENOME_BP,
        },
        'pairs': [{k: r[k] for k in ('human', 'banana', 'rel_type', 'human_length', 'aligned_residues', 'identical_residues',
                                     'aligned_nucleotides', 'identical_nucleotides')} for r in results],
    }


if __name__ == '__main__':
    if '--fetch' in sys.argv:
        retrieve()
    snapshot, hashes = load()
    result = analyze(snapshot)
    result['file_sha256'] = hashes
    pairs = result.pop('pairs')
    # Readable identifiers from the gene-count snapshot (UniProt entry names or RefSeq/Ensembl IDs).
    genes = json.loads(gzip.decompress(GENE_SNAPSHOT.read_bytes()))
    canonical = {row[0]: row[1] for code in ('HUMAN', 'MUSAM') for row in genes['proteins'][code]}
    for row in pairs:
        row['human_canonical_id'] = canonical[row['human']]
        row['banana_canonical_id'] = canonical[row['banana']]
    (DIRECTORY / 'sequence_summary.json').write_text(json.dumps(result, indent=2) + '\n')
    with open(DIRECTORY / 'banana_pairs.csv', 'w', newline='\n') as handle:
        handle.write(','.join(pairs[0]) + '\n')
        for row in pairs:
            handle.write(','.join(str(v) for v in row.values()) + '\n')
    print(json.dumps(result, indent=2))