"""
Structure Prediction Engine — PDB Generator
Generates 3D backbone coordinates from Fano-constrained geometry.

Uses Alpha-Flex manifold data to derive phi/psi backbone angles,
then builds an all-atom backbone (N, CA, C, O) using standard
peptide bond geometry with NeRF (Natural Extension Reference Frame).
"""

import math
import numpy as np
from alphaflex_engine import engine, ALPHABET_SCALES, PHI

# Standard peptide bond geometry (Angstroms)
BOND_N_CA = 1.458
BOND_CA_C = 1.525
BOND_C_N = 1.329
BOND_C_O = 1.231

# Bond angles (radians)
ANG_N_CA_C = math.radians(111.2)
ANG_CA_C_N = math.radians(116.2)
ANG_C_N_CA = math.radians(121.7)
ANG_CA_C_O = math.radians(120.8)

AA_3 = {
    'A': 'ALA', 'C': 'CYS', 'D': 'ASP', 'E': 'GLU', 'F': 'PHE',
    'G': 'GLY', 'H': 'HIS', 'I': 'ILE', 'K': 'LYS', 'L': 'LEU',
    'M': 'MET', 'N': 'ASN', 'P': 'PRO', 'Q': 'GLN', 'R': 'ARG',
    'S': 'SER', 'T': 'THR', 'V': 'VAL', 'W': 'TRP', 'Y': 'TYR',
}


def _nerf(a, b, c, bond, angle, dihedral):
    """
    Natural Extension Reference Frame — place atom D given A, B, C,
    bond length C-D, angle B-C-D, and dihedral A-B-C-D.
    """
    bc = c - b
    bc = bc / (np.linalg.norm(bc) + 1e-15)
    ab = b - a
    n = np.cross(ab, bc)
    n_norm = np.linalg.norm(n)
    if n_norm < 1e-10:
        if abs(bc[0]) < 0.9:
            n = np.cross(bc, np.array([1, 0, 0]))
        else:
            n = np.cross(bc, np.array([0, 1, 0]))
    n = n / (np.linalg.norm(n) + 1e-15)
    m = np.cross(n, bc)

    d = np.array([
        -bond * math.cos(math.pi - angle),
        bond * math.sin(math.pi - angle) * math.cos(dihedral),
        bond * math.sin(math.pi - angle) * math.sin(dihedral),
    ])
    return c + d[0] * bc + d[1] * m + d[2] * n


def fano_to_phi_psi(sigma, det, lane, flex=1.0):
    """Convert Fano manifold metrics to backbone phi/psi angles (radians)."""
    templates = {
        0: (-57.8, -47.0),   # Alpha helix
        1: (-119.0, 113.0),  # Beta sheet parallel
        2: (-139.0, 135.0),  # Beta sheet anti-parallel
        3: (-57.8, -47.0),   # Alpha helix
        4: (-78.0, 149.0),   # Polyproline II
        5: (-66.0, -24.0),   # 3-10 helix
        6: (-60.0, -30.0),   # Pi helix
    }
    base_phi, base_psi = templates.get(lane % 7, (-57.8, -47.0))
    sig_pert = math.tanh(sigma * 2.0) * 15.0
    det_pert = math.tanh(det * 0.5) * 10.0
    phi = max(-180, min(180, base_phi + det_pert * math.cos(lane * PHI)))
    psi = max(-180, min(180, base_psi + sig_pert * math.sin(lane * PHI)))
    return math.radians(phi), math.radians(psi)


def _ss_from_lane(lane):
    """Derive secondary structure label from Fano lane."""
    if lane in (0, 3):
        return 'H'  # Helix lanes
    elif lane in (1, 2):
        return 'E'  # Sheet lanes
    elif lane == 4:
        return 'C'  # PPII / coil
    elif lane == 5:
        return 'H'  # 3-10 helix
    elif lane == 6:
        return 'H'  # Pi helix
    return 'C'


def predict_structure(sequence, biases=None, flex=1.0):
    """
    Predict 3D backbone from sequence using Fano-constrained angles.
    Returns PDB string, atom list, per-residue metrics.
    """
    if biases is None:
        biases = np.ones(8)
    else:
        biases = np.array(biases, dtype=np.float64)

    valid = set(ALPHABET_SCALES.keys())
    seq = "".join(c for c in sequence.upper() if c in valid)
    if len(seq) < 2:
        return {"error": "Sequence too short (need >= 2 residues)"}

    res = engine.run_chain(seq, biases, flex)
    sigs, dets, lanes = res['sigs'], res['dets'], res['fano_lanes']
    n = len(seq)

    # Compute phi/psi from Fano lane templates + sigma/det perturbation
    phis, psis = [], []
    ss = []
    for i in range(n):
        ph, ps = fano_to_phi_psi(sigs[i], dets[i], lanes[i], flex)
        phis.append(ph)
        psis.append(ps)
        ss.append(_ss_from_lane(lanes[i]))

    # Build backbone with NeRF
    atoms = []
    N = np.array([0.0, 0.0, 0.0])
    CA = np.array([BOND_N_CA, 0.0, 0.0])
    C = _nerf(np.array([-1.0, 0.0, 0.0]), N, CA, BOND_CA_C, ANG_N_CA_C, psis[0])
    O = _nerf(N, CA, C, BOND_C_O, ANG_CA_C_O, math.pi)

    atoms += [("N", seq[0], 1, N), ("CA", seq[0], 1, CA), ("C", seq[0], 1, C), ("O", seq[0], 1, O)]

    prev_N, prev_CA, prev_C = N, CA, C

    for i in range(1, n):
        omega = math.pi  # Trans
        new_N = _nerf(prev_N, prev_CA, prev_C, BOND_C_N, ANG_CA_C_N, omega)
        new_CA = _nerf(prev_CA, prev_C, new_N, BOND_N_CA, ANG_C_N_CA, phis[i])
        new_C = _nerf(prev_C, new_N, new_CA, BOND_CA_C, ANG_N_CA_C, psis[i])
        new_O = _nerf(new_N, new_CA, new_C, BOND_C_O, ANG_CA_C_O, math.pi)

        rid = i + 1
        atoms += [("N", seq[i], rid, new_N), ("CA", seq[i], rid, new_CA),
                  ("C", seq[i], rid, new_C), ("O", seq[i], rid, new_O)]
        prev_N, prev_CA, prev_C = new_N, new_CA, new_C

    # PDB format
    ss_string = "".join(ss)
    lines = [
        "REMARK   1 LIGAND CONNECTIONS - FANO GEOMETRIC STRUCTURE PREDICTION",
        f"REMARK   2 SEQUENCE: {seq[:60]}{'...' if n > 60 else ''}",
        f"REMARK   3 LENGTH: {n} residues",
        f"REMARK   4 FLEX (OMEGA_H): {flex}",
        f"REMARK   5 MEAN_DET: {res['mean_det']:.6f}  MEAN_SIG: {res['mean_sig']:.6f}",
        f"REMARK   6 SS: {ss_string}",
        "REMARK   7 METHOD: Fano-constrained NeRF backbone generation",
    ]
    for idx, (name, aa, rid, pos) in enumerate(atoms):
        rn = AA_3.get(aa, 'UNK')
        el = name[0]
        bfactor = dets[min(rid - 1, n - 1)]
        lines.append(
            f"ATOM  {idx+1:5d} {name:<4s} {rn:>3s} A{rid:4d}    "
            f"{pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}"
            f"  1.00{bfactor:6.2f}           {el:>2s}"
        )
    lines.append("END")
    pdb_string = "\n".join(lines)

    metrics = [{
        "residue": seq[i], "position": i + 1,
        "ss": ss[i],
        "phi_deg": round(math.degrees(phis[i]), 1),
        "psi_deg": round(math.degrees(psis[i]), 1),
        "sigma": round(sigs[i], 6), "determinant": round(dets[i], 6),
        "fano_lane": lanes[i],
    } for i in range(n)]

    return {
        "sequence": seq, "length": n, "atom_count": len(atoms),
        "mean_det": round(res['mean_det'], 6),
        "mean_sig": round(res['mean_sig'], 6),
        "secondary_structure": ss_string,
        "ss_counts": {"helix": ss.count('H'), "sheet": ss.count('E'), "coil": ss.count('C')},
        "pdb_string": pdb_string, "metrics": metrics,
        "phi_angles": [round(math.degrees(p), 1) for p in phis],
        "psi_angles": [round(math.degrees(p), 1) for p in psis],
        "fano_lanes": lanes,
    }


def batch_screen(sequences, biases=None, flex=1.0):
    """Screen multiple sequences, rank by Fano binding fitness."""
    if biases is None:
        biases = np.ones(8)
    else:
        biases = np.array(biases, dtype=np.float64)

    valid_aa = set(ALPHABET_SCALES.keys())
    results = []

    for i, raw in enumerate(sequences):
        seq = "".join(c for c in raw.upper().strip() if c in valid_aa)
        if len(seq) < 2:
            continue
        r = engine.run_chain(seq, biases, flex)
        fitness = r['mean_det'] * math.exp(-r['mean_sig'])

        ar = {'W', 'Y', 'F'}
        stacks = sum(1 for j in range(len(seq) - 4)
                     if seq[j] in ar and (seq[j+3] in ar or seq[j+4] in ar))
        fitness *= (1.0 + 0.25 * stacks)

        results.append({
            "index": i, "sequence": seq, "length": len(seq),
            "fitness": round(fitness, 8),
            "mean_det": round(r['mean_det'], 6),
            "mean_sig": round(r['mean_sig'], 6),
            "max_det": round(r['max_det'], 6),
            "min_sig": round(r['min_sig'], 6),
            "aromatic_stacks": stacks,
            "unique_fano_lanes": len(set(r['fano_lanes'])),
            "fano_lane_distribution": [r['fano_lanes'].count(k) for k in range(7)],
        })

    results.sort(key=lambda x: x['fitness'], reverse=True)
    for rank, r in enumerate(results):
        r['rank'] = rank + 1
    return results
