"""Protein function prediction using a simplified GCN-like approach.

This is a lightweight approximation inspired by DeepFRI. For production use,
bake the full DeepFRI weights into the Docker image (see Dockerfile additions).

Outputs:
- GO term predictions with confidence scores (MF / BP / CC)
- Per-residue importance scores (saliency map)
- Per-residue amino acid composition

Note: EC number prediction is not implemented in the heuristic model.
"""

from __future__ import annotations

import json
import logging
import urllib.request

logger = logging.getLogger(__name__)


def _fetch_pdb_sequence(pdb_id: str) -> str:
    """Fetch the amino acid sequence for a PDB entry from RCSB."""
    url = f"https://data.rcsb.org/rest/v1/core/polymer_entity/{pdb_id}/1"
    try:
        data = json.loads(urllib.request.urlopen(url, timeout=15).read())
        return data.get("entity_poly", {}).get("pdbx_seq_one_letter_code_can", "")
    except Exception:
        pass

    # Fallback: fetch FASTA
    try:
        url = f"https://www.rcsb.org/fasta/entry/{pdb_id}"
        text = urllib.request.urlopen(url, timeout=15).read().decode()
        lines = [l for l in text.splitlines() if not l.startswith(">")]
        return "".join(lines).replace("\n", "")
    except Exception as e:
        raise RuntimeError(f"Could not fetch sequence for {pdb_id}: {e}")


def _predict_from_sequence(sequence: str, pdb_id: str) -> dict:
    """Lightweight function prediction based on sequence composition.

    This is a heuristic approximation. Replace with proper GCN inference
    when DeepFRI weights are baked into the Docker image.
    """
    seq_upper = sequence.upper()
    seq_len = len(seq_upper)
    aa_comp = {}
    for aa in seq_upper:
        aa_comp[aa] = aa_comp.get(aa, 0) + 1

    # Predict GO terms based on amino acid composition patterns
    predicted_go = []
    confidence_base = 0.5

    # Simple composition-based predictions
    hydrophobic_fraction = sum(aa_comp.get(a, 0) for a in "AILMFWV") / max(seq_len, 1)
    charged_fraction = sum(aa_comp.get(a, 0) for a in "DEKRH") / max(seq_len, 1)

    if hydrophobic_fraction > 0.4:
        predicted_go.append({
            "go_id": "GO:0016020",
            "name": "membrane",
            "namespace": "CC",
            "confidence": round(min(0.6 + hydrophobic_fraction * 0.3, 0.95), 3),
        })
    if charged_fraction > 0.25:
        predicted_go.append({
            "go_id": "GO:0005515",
            "name": "protein binding",
            "namespace": "MF",
            "confidence": round(min(0.55 + charged_fraction * 0.2, 0.9), 3),
        })
    if hydrophobic_fraction > 0.4 and charged_fraction > 0.15:
        predicted_go.append({
            "go_id": "GO:0007165",
            "name": "signal transduction",
            "namespace": "BP",
            "confidence": round(min(0.5 + hydrophobic_fraction * 0.2, 0.85), 3),
        })

    # Always include a general prediction
    predicted_go.append({
        "go_id": "GO:0003674",
        "name": "molecular_function",
        "namespace": "MF",
        "confidence": 0.99,
    })

    # Per-residue importance (saliency approximation)
    # Higher importance at charged/polar residues on the surface
    saliency = []
    for i, aa in enumerate(seq_upper):
        score = 0.1
        if aa in "DEKRH":
            score = 0.6
        elif aa in "STNQ":
            score = 0.4
        elif aa in "AGV":
            score = 0.2
        else:
            score = 0.15
        saliency.append(round(score, 3))

    # Actual per-residue composition (fractions), sorted most abundant first
    composition = {
        "aa": "ACDEFGHIKLMNPQRSTVWY",
        "fractions": {aa: round(aa_comp.get(aa, 0) / max(seq_len, 1), 4) for aa in "ACDEFGHIKLMNPQRSTVWY"},
    }

    return {
        "pdb_id": pdb_id.upper(),
        "sequence_length": seq_len,
        "go_terms": predicted_go,
        "ec_numbers": [],
        "saliency": saliency,
        "composition": composition,
        "method": "heuristic_composition",
        "note": "Predictions based on amino acid composition. EC number prediction is not implemented in this heuristic model; for research-grade predictions use the full DeepFRI model.",
    }


def predict_function(pdb_id: str) -> dict:
    """Main entry point: predict protein function from structure."""
    sequence = _fetch_pdb_sequence(pdb_id)
    if not sequence:
        raise RuntimeError(f"No sequence available for PDB {pdb_id}")
    return _predict_from_sequence(sequence, pdb_id)