
import numpy as np
import biotite.structure as struc
from typing import Dict, List, Tuple
import logging

logger = logging.getLogger(__name__)

# --- Optional Numba JIT Support ---
try:
    from numba import njit
except ImportError:
    def njit(func=None, **kwargs):
        if func is None:
            return lambda f: f
        return func

# --- Random Coil Chemical Shifts (Wishart et al.) ---
# EDUCATIONAL NOTE - Random Coil Shifts:
# ======================================
# "Random Coil" refers to a protein state with no fixed secondary structure (a flexible chain).
# The chemical shift of an atom in a random coil depends primarily on its amino acid type.
#
# These values serve as the "baseline" or "zero point" for structure prediction.
# Any deviation from these values (Secondary Shift) indicates structural formation:
# - Alpha Helix formation moves C-alpha downfield (higher ppm) and N upfield (lower ppm).
# - Beta Sheet formation moves C-alpha upfield (lower ppm) and N downfield (higher ppm).
#
# Reference: Wishart, D.S. et al. (1995) J. Biomol. NMR.
# Referenced to DSS at 25C.
# Values for: HA, CA, CB, C, N, HN (Amide H)
# Units: ppm
RANDOM_COIL_SHIFTS: Dict[str, Dict[str, float]] = {
    "ALA": {"HA": 4.32, "CA": 52.5, "CB": 19.1, "C": 177.8, "N": 123.8, "H": 8.24},
    "ARG": {"HA": 4.34, "CA": 56.0, "CB": 30.9, "C": 176.3, "N": 121.3, "H": 8.23},
    "ASN": {"HA": 4.75, "CA": 53.1, "CB": 38.9, "C": 175.2, "N": 118.7, "H": 8.75},
    "ASP": {"HA": 4.66, "CA": 54.2, "CB": 41.1, "C": 176.3, "N": 120.4, "H": 8.34},
    "CYS": {"HA": 4.69, "CA": 58.2, "CB": 28.0, "C": 174.6, "N": 118.8, "H": 8.32},
    "CYX": {"HA": 4.69, "CA": 58.2, "CB": 28.0, "C": 174.6, "N": 118.8, "H": 8.32},
    "GLN": {"HA": 4.32, "CA": 56.0, "CB": 29.4, "C": 176.0, "N": 120.4, "H": 8.25},
    "GLU": {"HA": 4.29, "CA": 56.6, "CB": 29.9, "C": 176.6, "N": 120.2, "H": 8.35},
    "GLY": {"HA": 3.96, "CA": 45.1, "CB": 0.0,  "C": 174.9, "N": 108.8, "H": 8.33},
    "HIS": {"HA": 4.63, "CA": 55.0, "CB": 29.0, "C": 174.1, "N": 118.2, "H": 8.42},
    "HID": {"HA": 4.63, "CA": 55.0, "CB": 29.0, "C": 174.1, "N": 118.2, "H": 8.42},
    "HIE": {"HA": 4.63, "CA": 55.0, "CB": 29.0, "C": 174.1, "N": 118.2, "H": 8.42},
    "HIP": {"HA": 4.63, "CA": 55.0, "CB": 29.0, "C": 174.1, "N": 118.2, "H": 8.42},
    "ILE": {"HA": 4.17, "CA": 61.1, "CB": 38.8, "C": 176.4, "N": 121.4, "H": 8.00},
    "LEU": {"HA": 4.34, "CA": 55.1, "CB": 42.4, "C": 177.6, "N": 121.8, "H": 8.16},
    "LYS": {"HA": 4.32, "CA": 56.2, "CB": 33.1, "C": 176.6, "N": 120.4, "H": 8.29},
    "MET": {"HA": 4.48, "CA": 55.4, "CB": 32.6, "C": 176.3, "N": 119.6, "H": 8.28},
    "PHE": {"HA": 4.62, "CA": 57.7, "CB": 39.6, "C": 175.8, "N": 120.3, "H": 8.12},
    "PRO": {"HA": 4.42, "CA": 63.3, "CB": 32.1, "C": 177.3, "N": 0.0,   "H": 0.0},  # No Amide N/H
    "SER": {"HA": 4.47, "CA": 58.3, "CB": 63.8, "C": 174.6, "N": 115.7, "H": 8.31},
    "THR": {"HA": 4.35, "CA": 61.8, "CB": 69.8, "C": 174.7, "N": 113.6, "H": 8.15},
    "TRP": {"HA": 4.66, "CA": 57.5, "CB": 29.6, "C": 176.1, "N": 121.3, "H": 8.25},
    "TYR": {"HA": 4.55, "CA": 57.9, "CB": 38.8, "C": 175.9, "N": 120.3, "H": 8.12},
    "VAL": {"HA": 4.12, "CA": 62.2, "CB": 32.9, "C": 176.3, "N": 119.9, "H": 8.03},
}

# Module-level variable for noise scale, allowing easy monkeypatching in tests
_NOISE_SCALE = 0.15


# --- Secondary Structure Offsets (SPARTA+) ---
# EDUCATIONAL NOTE - Secondary Chemical Shifts:
# =============================================
# The local magnetic field experienced by a nucleus is heavily influenced by the
# geometry of the protein backbone (Phi/Psi angles).
#
# SPARTA+ (Shift Prediction from Analogy in Residue-type and Torsion Angle):
# It predicts chemical shifts by finding homologous structures with similar geometry.
#
# Our implementation uses simple statistical offsets instead of database mining,
# but follows the same principle: Geometry determines Shift.
#
# Reference State: DSS (4,4-dimethyl-4-silapentane-1-sulfonic acid)
# This is the "Zero" for proton/carbon NMR, much like sea level for altitude.
# Using a standard reference ensures shifts are comparable across different labs.
#
# Approximate mean offsets for Helical and Sheet conformations relative to random coil
# Based on general statistics (e.g. Spera & Bax 1991)
# Format: {metric: {Helix: val, Sheet: val}}
SECONDARY_SHIFTS: Dict[str, Dict[str, float]] = {
    # C-alpha: Shifted downfield (positive) in Helix, upfield (negative) in Sheet
    "CA": {"alpha": 3.1,  "beta": -1.5},
    # C-beta: Opposite trend to C-alpha
    "CB": {"alpha": -0.5, "beta": 2.2},
    # Carbonyl Carbon: Follows C-alpha trend
    "C":  {"alpha": 2.2,  "beta": -1.6},
    # H-alpha: Shifted upfield (negative) in Helix, downfield (positive) in Sheet
    "HA": {"alpha": -0.4, "beta": 0.5},
    # Amide N: Complex, but generally upfield in Helix
    "N":  {"alpha": -1.5, "beta": 1.2},
    "H":  {"alpha": -0.2, "beta": 0.3},
}

# --- Ring Current Intensity Factors ---
# EDUCATIONAL NOTE - Ring Current Physics:
# ========================================
# Aromatic rings (Benchmark: Benzene) have delocalized pi-electrons that circulate
# when exposed to a magnetic field, creating an opposing induced magnetic field.
#
# - Regions ABOVE/BELOW the ring are SHIELDED (Field opposes external field -> Lower ppm).
# - Regions in the PLANE of the ring are DESHIELDED (Field adds to external field -> Higher ppm).
#
# Model: Point Dipole approximation.
# Shift = Intensity * (1 - 3*cos^2(theta)) / r^3
#
# References for further reading:
# 1. Haigh, C. W., & Mallion, R. B. (1980). "Ring current theories in nuclear magnetic resonance". 
#    Progress in Nuclear Magnetic Resonance Spectroscopy, 13(4), 303-344.
# 2. Pople, J. A. (1956). "Proton magnetic shielding in aromatic compounds". 
#    The Journal of Chemical Physics, 24(5), 1111.
# 3. Case, D. A. (1995). "Chemical shifts in proteins". 
#    Current Opinion in Structural Biology, 5(2), 272-276.
#
# Intensities are relative to Benzene.
RING_INTENSITIES = {
    "PHE": 1.2,  # Benzene ring (Standard)
    "TYR": 1.2,  # Phenol ring (Similar to Benzene)
    "TRP": 1.3,  # Indole (Stronger system)
    "HIS": 0.5,  # Imidazole (Weaker, depends on protonation)
    "HID": 0.5,
    "HIE": 0.5,
    "HIP": 0.5,
}


from synth_nmr.structure_utils import get_secondary_structure

def predict_chemical_shifts(structure: struc.AtomArray) -> Dict[str, Dict[str, Dict[str, float]]]:
    """
    Predict chemical shifts based on secondary structure and ring currents.
    
    This function combines random coil shifts, secondary structure-based offsets (SPARTA+-like),
    and ring current effects from aromatic residues to predict protein chemical shifts.

    EDUCATIONAL NOTE - Prediction Algorithm:
    ========================================
    1. Calculate Backbone Dihedrals (Phi/Psi) for every residue.
    2. Classify Secondary Structure:
       - Alpha: Phi ~ -60, Psi ~ -45
       - Beta:  Phi ~ -120, Psi ~ 120
       - Coil:  Everything else
    3. Calculate Shift:
       Shift = Random_Coil + Structure_Offset + Noise
       
    LIMITATIONS:
    - Ring Current Effects (for protons): While an O(N^2) geometry check was previously
      a concern, these effects are now included for protons near aromatic rings
      (Phe, Tyr, Trp, His) using a point-dipole approximation. This is crucial for
      protons in close proximity to aromatic systems. Carbon atoms are not currently
      included for ring current effects.
    - H-Bonding: Hydrogen bonds affect Amide H shifts significantly. We omit this for simplicity.
    - Sequence History: Real shifts depend on (i-1) and (i+1) neighbor types. We omit this for simplicity.
    
    Args:
        structure: A biotite.structure.AtomArray containing the protein. Must not be empty.
        
    Returns:
        A nested dictionary of predicted shifts: {chain_id: {res_id: {atom_name: value}}}
        
    Raises:
        TypeError: If the input is not a biotite.structure.AtomArray.
        ValueError: If the input structure is empty.
    """
    logger.info("Predicting chemical shifts (SPARTA+ model with ring currents)...")

    # 1. Input Validation
    if not isinstance(structure, struc.AtomArray):
        raise TypeError("Input 'structure' must be a biotite.structure.AtomArray.")
    if structure.array_length() == 0:
        logger.warning("Input 'structure' is empty. Cannot predict chemical shifts.")
        return {}

    try:
        # 2. Get Secondary Structure and Aromatic Ring Info
        ss_list = get_secondary_structure(structure)
        rings = _get_aromatic_rings(structure)
        if rings.size > 0:
            logger.debug(f"Found {rings.shape[0]} aromatic rings for ring current calculation.")

        # 3. Iterate through residues and calculate shifts
        results = {}
        for i, start_idx in enumerate(struc.get_residue_starts(structure)):
            end_idx = struc.get_residue_starts(structure)[i+1] if i + 1 < len(struc.get_residue_starts(structure)) else None
            res_atoms = structure[start_idx:end_idx]
            
            res_id = res_atoms.res_id[0]
            res_name = res_atoms.res_name[0]
            chain_id = res_atoms.chain_id[0]

            if res_name not in RANDOM_COIL_SHIFTS:
                logger.debug(f"Skipping non-standard residue: {res_name} {res_id}")
                continue

            ss_state = ss_list[i] if i < len(ss_list) else "coil"
            logger.debug(f"Processing ResID {res_id} ({res_name}), SS: {ss_state}")

            rc_shifts = RANDOM_COIL_SHIFTS[res_name]
            atom_shifts = {}

            for atom_type, base_val in rc_shifts.items():
                if base_val == 0:  # Skip atoms with no defined random coil shift (e.g., Proline H)
                    continue

                # Start with base random coil value
                val = base_val
                
                # Add secondary structure offset
                offset = SECONDARY_SHIFTS.get(atom_type, {}).get(ss_state, 0.0)
                val += offset
                
                # Add ring current shift for protons
                if rings.size > 0 and atom_type.startswith('H'):
                    try:
                        target_atom = res_atoms[res_atoms.atom_name == atom_type][0]
                        rc_shift = _calculate_ring_current_shift(target_atom.coord, rings)
                        val += rc_shift
                    except IndexError:
                        # Atom not found in this specific residue, skip ring current
                        logger.debug(f"Atom {atom_type} not found in residue {res_id} for ring current calculation.")
                        pass

                # Add small random noise for "realism" (0.1 - 0.3 ppm)
                # Experimental assignments always have error/variation
                noise = np.random.normal(0, _NOISE_SCALE) if base_val != 0 else 0
                val += noise

                atom_shifts[atom_type] = round(val, 3)

            if chain_id not in results:
                results[chain_id] = {}
            results[chain_id][res_id] = atom_shifts
            
        logger.info(f"Successfully predicted chemical shifts for {struc.get_residue_count(structure)} residues.")
        return results

    except Exception as e:
        logger.error(f"An unexpected error occurred during chemical shift prediction: {e}", exc_info=True)
        raise

def calculate_csi(
    shifts: Dict[str, Dict[int, Dict[str, float]]], 
    structure: struc.AtomArray
) -> Dict[str, Dict[int, float]]:
    """
    Calculate the Chemical Shift Index (CSI) for C-alpha atoms.
    
    The CSI is the deviation of an observed chemical shift from its random coil value.
    It is a reliable indicator of secondary structure.
    - Positive Delta(CA) > 0.7 ppm suggests a Helical conformation.
    - Negative Delta(CA) < -0.7 ppm suggests a Sheet conformation.
    
    Args:
        shifts: A dictionary of chemical shifts, as produced by `predict_chemical_shifts`.
        structure: A biotite.structure.AtomArray, required for mapping residue IDs to names.
        
    Returns:
        A dictionary containing the C-alpha CSI for each residue: {chain_id: {res_id: delta_ppm}}
        
    Raises:
        TypeError: If inputs are not of the correct type.
        ValueError: If inputs are empty.
    """
    logger.info("Calculating Chemical Shift Index (CSI) for C-alpha atoms...")

    # 1. Input Validation
    if not isinstance(shifts, dict):
        raise TypeError("Input 'shifts' must be a dictionary.")
    if not shifts:
        logger.warning("Input 'shifts' dictionary is empty. Returning no CSI data.")
        return {}
    if not isinstance(structure, struc.AtomArray):
        raise TypeError("Input 'structure' must be a biotite.structure.AtomArray.")
    if structure.array_length() == 0:
        logger.warning("Input 'structure' is empty. Cannot calculate CSI.")
        return {}

    try:
        # 2. Create a mapping from residue ID to residue name for quick lookup
        res_names = {}
        for idx in struc.get_residue_starts(structure):
            res = structure[idx]
            res_names[res.res_id] = res.res_name
        if not res_names:
            logger.warning("Could not create a residue map from the provided structure.")
            return {}

        # 3. Calculate CSI
        csi_data = {}
        for chain_id, chain_shifts in shifts.items():
            if not isinstance(chain_shifts, dict): continue
            csi_data[chain_id] = {}
            
            for res_id, atom_shifts in chain_shifts.items():
                if not isinstance(atom_shifts, dict): continue

                res_name = res_names.get(res_id)
                if not res_name:
                    logger.debug(f"Residue ID {res_id} from shifts not found in structure. Skipping.")
                    continue

                if "CA" in atom_shifts and res_name in RANDOM_COIL_SHIFTS:
                    measured = atom_shifts["CA"]
                    random_coil_val = RANDOM_COIL_SHIFTS[res_name].get("CA")
                    
                    if random_coil_val is not None:
                        delta = measured - random_coil_val
                        csi_data[chain_id][res_id] = round(delta, 3)
                        logger.debug(f"CSI for {res_name} {res_id}: Measured={measured}, RC={random_coil_val}, Delta={delta}")
                    else:
                        logger.debug(f"No random coil 'CA' value for {res_name}. Skipping CSI calculation for this residue.")
        
        logger.info("CSI calculation complete.")
        return csi_data

    except Exception as e:
        logger.error(f"An unexpected error occurred during CSI calculation: {e}", exc_info=True)
        raise

def _get_aromatic_rings(structure):
    """
    Identify aromatic rings and calculate their centers and normal vectors.
    """
    rings = []
    
    # Iterate residues
    res_starts = struc.get_residue_starts(structure)
    for idx in res_starts:
        res = structure[idx]
        res_name = res.res_name
        
        if res_name in RING_INTENSITIES:
            # Extract ring atoms to calculate geometry
            # Simplified definition of ring atoms
            res_slice = structure[structure.res_id == res.res_id]
            
            if res_name in ["PHE", "TYR"]:
                # 6-membered ring: CG, CD1, CD2, CE1, CE2, CZ
                ring_atoms = res_slice[np.isin(res_slice.atom_name, ["CG", "CD1", "CD2", "CE1", "CE2", "CZ"])]
            elif res_name == "TRP":
                # Indole is 9 atoms, effective center near CD2/CE2 bond.
                # Simplified: averaging all ring atoms
                ring_names = ["CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"]
                ring_atoms = res_slice[np.isin(res_slice.atom_name, ring_names)]
            elif res_name in ["HIS", "HID", "HIE", "HIP"]:
                # 5-membered ring: CG, ND1, CD2, CE1, NE2
                ring_atoms = res_slice[np.isin(res_slice.atom_name, ["CG", "ND1", "CD2", "CE1", "NE2"])]
            else:
                continue
                
            if len(ring_atoms) >= 3:
                # Geometric Center
                center = np.mean(ring_atoms.coord, axis=0)
                
                # Normal Vector (Cross product of two vectors in the ring)
                # v1: Center -> Atom 0
                # v2: Center -> Atom 1
                # Normal = v1 x v2
                v1 = ring_atoms[0].coord - center
                v2 = ring_atoms[1].coord - center
                normal = np.cross(v1, v2)
                norm = np.linalg.norm(normal)
                if norm > 0:
                    normal = normal / norm
                    intensity = RING_INTENSITIES[res_name]
                    rings.append((center, normal, intensity))
                    
    if not rings:
        return np.empty((0, 7), dtype=np.float64)
        
    # Convert list of tuples (center, normal, intensity) to (N, 7) array
    ring_array = np.zeros((len(rings), 7), dtype=np.float64)
    for i, (c, n, intensity) in enumerate(rings):
        ring_array[i, 0:3] = c
        ring_array[i, 3:6] = n
        ring_array[i, 6] = intensity
        
    return ring_array

@njit
def _calculate_ring_current_shift(proton_coord, rings):
    """
    Calculate total ring current shift for a proton from all rings.
    'rings' is a numpy array of shape (N, 7): [cx, cy, cz, nx, ny, nz, intensity]
    Formula: delta = Intensity * B_factor * (1 - 3*cos^2(theta)) / r^3
    """
    total_shift = 0.0
    # B_FACTOR is an empirical scaling constant, typically derived from fitting
    # experimental data or theoretical calculations for a reference aromatic system.
    # It converts the dimensionless geometric factor and intensity into ppm.
    # Typical values range from 8 to 15 ppm*A^3.
    B_FACTOR = 11.0 # Empirical scaling factor (ppm * A^3)
    
    for j in range(rings.shape[0]):
        center = rings[j, 0:3]
        normal = rings[j, 3:6]
        intensity = rings[j, 6]
        
        # Vector from ring center to proton
        v = (proton_coord - center).astype(np.float64)
        r = np.sqrt(np.sum(v**2))
        
        if r < 1.0: 
            continue # Too close/clashing, ignore singularity
            
        # Cos(theta) = dot(v, n) / (|v|*|n|) -> |n|=1
        costheta = np.sum(v * normal) / r
        
        # Geometric Factor G(r, theta) = (1 - 3*cos^2(theta)) / r^3
        # If theta = 0 (above ring), cos=1 -> (1-3)/r^3 = -2/r^3 (Shielding)
        # If theta = 90 (in plane), cos=0 -> (1-0)/r^3 =  1/r^3 (Deshielding)
        geom_factor = (1.0 - 3.0 * costheta**2) / (r**3)
        
        shift = intensity * B_FACTOR * geom_factor
        
        total_shift += shift
        
    return total_shift
