import logging
from typing import List, Tuple

import numpy as np

from synth_pdb.orientogram import compute_6d_orientations
from synth_pdb.validator import RAMACHANDRAN_POLYGONS, PDBValidator

logger = logging.getLogger(__name__)

# Re-export core function for Latent Space Explorer
__all__ = ["compute_6d_orientations", "extract_quality_features", "get_feature_names"]


def get_feature_names() -> List[str]:
    """Returns the list of feature names in the order they appear in the feature vector."""
    return [
        "ramachandran_favored_pct",
        "ramachandran_outliers_pct",
        "clash_count",
        "bond_length_violation_count",
        "bond_angle_violation_count",
        "peptide_bond_violation_count",
        "radius_of_gyration",
        "mean_b_factor",
    ]


def extract_quality_features(pdb_content: str) -> np.ndarray:
    """
    Extracts a feature vector for the Protein Structure Quality Classifier.

    Features:
    0. Ramachandran Favored %
    1. Ramachandran Outliers %
    2. Steric Clash Count
    3. Bond Length Violation Count
    4. Bond Angle Violation Count
    5. Peptide Bond Planarity Violation Count
    6. Radius of Gyration (Compactness)
    7. Mean B-factor (Flexibility)

    Returns:
        np.ndarray: A 1D array of floats (shape: 8,)
    """
    try:
        # Use PDBValidator for geometric violations; intercept counts, not strings.
        validator = PDBValidator(pdb_content)

        # --- Ramachandran Analysis ---
        # We compute this manually to obtain percentages; the validator only reports outliers.
        phi, psi = _get_dihedrals(validator)
        rama_favored, rama_outliers = _analyze_ramachandran(phi, psi, validator)
        total_residues = len(phi)

        rama_favored_pct = (rama_favored / total_residues * 100) if total_residues > 0 else 0
        rama_outliers_pct = (rama_outliers / total_residues * 100) if total_residues > 0 else 0

        # --- Steric Clashes ---
        # backbone_only=True reduces noise from sidechain generation artefacts.
        initial_violations = len(validator.violations)
        validator.validate_steric_clashes(
            min_atom_distance=0.5, min_ca_distance=3.0, backbone_only=True
        )
        clash_count = len(validator.violations) - initial_violations

        # --- Bond Lengths (0.1 Å tolerance) ---
        initial_violations = len(validator.violations)
        validator.validate_bond_lengths(tolerance=0.1)
        bond_len_count = len(validator.violations) - initial_violations

        # --- Bond Angles (10° tolerance) ---
        initial_violations = len(validator.violations)
        validator.validate_bond_angles(tolerance=10.0)
        bond_ang_count = len(validator.violations) - initial_violations

        # --- Peptide Bond Planarity ---
        initial_violations = len(validator.violations)
        validator.validate_peptide_plane(tolerance_deg=30.0)
        peptide_plane_count = len(validator.violations) - initial_violations

        # --- Global Properties ---
        coords = np.array([atom["coords"] for atom in validator.atoms])
        b_factors = np.array([atom["temp_factor"] for atom in validator.atoms])

        # Number of residues used for per-residue normalisation of violation counts.
        num_residues: float = float(
            max(1, sum(len(res_dict) for res_dict in validator.grouped_atoms.values()))
        )

        # Radius of Gyration
        rg = 0.0
        if len(coords) > 0:
            center_of_mass = np.mean(coords, axis=0)
            rg = np.sqrt(np.sum(np.linalg.norm(coords - center_of_mass, axis=1) ** 2) / len(coords))

        mean_b_factor = np.mean(b_factors) if len(b_factors) > 0 else 0.0

        return np.array(
            [
                rama_favored_pct,
                rama_outliers_pct,
                float(clash_count) / num_residues,
                float(bond_len_count) / num_residues,
                float(bond_ang_count) / num_residues,
                float(peptide_plane_count) / num_residues,
                rg,
                mean_b_factor,
            ]
        )

    except Exception as e:
        logger.error(
            "extract_quality_features failed: %s. "
            "The caller is responsible for handling this exception.",
            e,
            exc_info=True,
        )
        raise


def _get_dihedrals(validator: PDBValidator) -> Tuple[List[float], List[float]]:
    """Extracts Phi/Psi angles from the validator's parsed atoms."""
    phi_list = []
    psi_list = []

    for _chain_id, residues_in_chain in validator.grouped_atoms.items():
        sorted_res_numbers = sorted(residues_in_chain.keys())
        for i, res_num in enumerate(sorted_res_numbers):
            current_res_atoms = residues_in_chain[res_num]

            # Phi
            phi = None
            if i > 0:
                prev_res_num = sorted_res_numbers[i - 1]
                prev_res_atoms = residues_in_chain.get(prev_res_num)
                if prev_res_atoms:
                    c_prev = prev_res_atoms.get("C")
                    n_curr = current_res_atoms.get("N")
                    ca_curr = current_res_atoms.get("CA")
                    c_curr = current_res_atoms.get("C")
                    if c_prev and n_curr and ca_curr and c_curr:
                        p1 = c_prev["coords"]
                        p2 = n_curr["coords"]
                        p3 = ca_curr["coords"]
                        p4 = c_curr["coords"]
                        phi = validator._calculate_dihedral_angle(p1, p2, p3, p4)

            # Psi
            psi = None
            if i < len(sorted_res_numbers) - 1:
                next_res_num = sorted_res_numbers[i + 1]
                next_res_atoms = residues_in_chain.get(next_res_num)
                if next_res_atoms:
                    n_curr = current_res_atoms.get("N")
                    ca_curr = current_res_atoms.get("CA")
                    c_curr = current_res_atoms.get("C")
                    n_next = next_res_atoms.get("N")
                    if n_curr and ca_curr and c_curr and n_next:
                        p1 = n_curr["coords"]
                        p2 = ca_curr["coords"]
                        p3 = c_curr["coords"]
                        p4 = n_next["coords"]
                        psi = validator._calculate_dihedral_angle(p1, p2, p3, p4)

            if phi is not None and psi is not None:
                phi_list.append(phi)
                psi_list.append(psi)

    return phi_list, psi_list


def _analyze_ramachandran(
    phi_list: List[float], psi_list: List[float], validator: PDBValidator
) -> Tuple[int, int]:
    """Counts favored and outlier residues."""
    favored = 0
    outliers = 0

    # We use "General" polygons for everything to simplify feature extraction
    # This keeps the model robust and noise-tolerant
    polygons = RAMACHANDRAN_POLYGONS["General"]

    for phi, psi in zip(phi_list, psi_list):
        is_favored = False
        for poly in polygons["Favored"]:
            if validator._is_point_in_polygon((phi, psi), poly):
                is_favored = True
                break

        if is_favored:
            favored += 1
        else:
            is_allowed = False
            for poly in polygons["Allowed"]:
                if validator._is_point_in_polygon((phi, psi), poly):
                    is_allowed = True
                    break

            if not is_allowed:
                outliers += 1

    return favored, outliers
