from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, List, Optional, Tuple

import numpy as np

_HELIX_PHI = (-100.0, -30.0)
_HELIX_PSI = (-80.0, -10.0)
_SHEET_PHI = (-180.0, -100.0)
_SHEET_PSI = (90.0, 180.0)


@dataclass(frozen=True)
class _ResidueBackbone:
    resname: str
    resid: int
    chain_id: int
    n_idx: Optional[int]
    ca_idx: Optional[int]
    c_idx: Optional[int]


def _all_resid_mask(system) -> str:
    atoms = system.atom_table()
    resids = atoms.get("resid", [])
    if not resids:
        return "resid 0:0"
    return f"resid {min(resids)}:{max(resids)}"


def _selection_set(selection) -> set:
    return {int(i) for i in selection.indices}


def _build_residues(system, selection_mask: Optional[Iterable[int]]) -> List[_ResidueBackbone]:
    atoms = system.atom_table()
    names = atoms.get("name", [])
    resnames = atoms.get("resname", [])
    resids = atoms.get("resid", [])
    chain_ids = atoms.get("chain_id", [])
    n_atoms = len(names)
    selection = set(selection_mask) if selection_mask else set(range(n_atoms))

    residues: List[_ResidueBackbone] = []
    current_key = None
    current_indices: List[int] = []

    def flush():
        if not current_indices:
            return
        if selection and not any(idx in selection for idx in current_indices):
            return
        n_idx = ca_idx = c_idx = None
        for idx in current_indices:
            name = str(names[idx]).strip().upper()
            if name == "N":
                n_idx = idx
            elif name == "CA":
                ca_idx = idx
            elif name == "C":
                c_idx = idx
        chain, resid, resname = current_key
        residues.append(
            _ResidueBackbone(
                resname=str(resname),
                resid=int(resid),
                chain_id=int(chain),
                n_idx=n_idx,
                ca_idx=ca_idx,
                c_idx=c_idx,
            )
        )

    for idx in range(n_atoms):
        key = (
            chain_ids[idx] if idx < len(chain_ids) else 0,
            resids[idx] if idx < len(resids) else 0,
            resnames[idx] if idx < len(resnames) else "RES",
        )
        if current_key is None:
            current_key = key
        if key != current_key:
            flush()
            current_indices = []
            current_key = key
        current_indices.append(idx)
    flush()
    return residues


def _dihedral(p0, p1, p2, p3) -> Optional[float]:
    b0 = p1 - p0
    b1 = p2 - p1
    b2 = p3 - p2
    norm_b1 = np.linalg.norm(b1)
    if norm_b1 == 0.0:
        return None
    b1 /= norm_b1
    v = b0 - np.dot(b0, b1) * b1
    w = b2 - np.dot(b2, b1) * b1
    norm_v = np.linalg.norm(v)
    norm_w = np.linalg.norm(w)
    if norm_v == 0.0 or norm_w == 0.0:
        return None
    v /= norm_v
    w /= norm_w
    x = np.dot(v, w)
    y = np.dot(np.cross(b1, v), w)
    return float(np.degrees(np.arctan2(y, x)))


def _classify(phi: Optional[float], psi: Optional[float]) -> str:
    if phi is None or psi is None:
        return "C"
    if _HELIX_PHI[0] <= phi <= _HELIX_PHI[1] and _HELIX_PSI[0] <= psi <= _HELIX_PSI[1]:
        return "H"
    if _SHEET_PHI[0] <= phi <= _SHEET_PHI[1] and _SHEET_PSI[0] <= psi <= _SHEET_PSI[1]:
        return "E"
    return "C"


def _phi_psi(residues: List[_ResidueBackbone], coords: np.ndarray, idx: int) -> Tuple[Optional[float], Optional[float]]:
    res = residues[idx]
    if res.n_idx is None or res.ca_idx is None or res.c_idx is None:
        return None, None
    phi = None
    psi = None
    if idx > 0:
        prev = residues[idx - 1]
        if prev.c_idx is not None:
            phi = _dihedral(
                coords[prev.c_idx],
                coords[res.n_idx],
                coords[res.ca_idx],
                coords[res.c_idx],
            )
    if idx + 1 < len(residues):
        nxt = residues[idx + 1]
        if nxt.n_idx is not None:
            psi = _dihedral(
                coords[res.n_idx],
                coords[res.ca_idx],
                coords[res.c_idx],
                coords[nxt.n_idx],
            )
    return phi, psi


def _avg_counts(ss: np.ndarray) -> dict:
    if ss.size == 0:
        return {}
    values, counts = np.unique(ss, return_counts=True)
    return {str(v): float(c) for v, c in zip(values, counts)}


def dssp(traj, system, mask: str = "protein", simplified: bool = False, chunk_frames: Optional[int] = None):
    """Return DSSP-like labels via phi/psi heuristic (H/E/C)."""
    selection = system.select(mask) if mask else system.select(_all_resid_mask(system))
    sel_set = _selection_set(selection)
    if not sel_set:
        selection = system.select(_all_resid_mask(system))
        sel_set = _selection_set(selection)
    residues = _build_residues(system, sel_set)
    labels = np.asarray([f"{r.resname}:{r.resid}" for r in residues], dtype="U16")
    n_res = len(residues)
    if n_res == 0:
        return labels, np.empty((0, 0), dtype="U1"), {}

    max_chunk = chunk_frames or 128
    ss_frames = []
    chunk = traj.read_chunk(max_chunk)
    if chunk is None:
        return labels, np.empty((0, n_res), dtype="U1"), {}
    while chunk is not None:
        coords = np.asarray(chunk["coords"], dtype=np.float64)
        for f in range(coords.shape[0]):
            frame_ss = np.empty(n_res, dtype="U1")
            for i in range(n_res):
                phi, psi = _phi_psi(residues, coords[f], i)
                frame_ss[i] = _classify(phi, psi)
            ss_frames.append(frame_ss)
        chunk = traj.read_chunk(max_chunk)

    if not ss_frames:
        return labels, np.empty((0, n_res), dtype="U1"), {}
    ss = np.stack(ss_frames, axis=0)
    avg = _avg_counts(ss)
    if simplified:
        return labels, ss, avg
    return labels, ss, avg


def dssp_allatoms(traj, system, mask: str = "protein", simplified: bool = False):
    """Return DSSP-like labels for all atoms."""
    return dssp(traj, system, mask=mask, simplified=simplified)[1]


def dssp_allresidues(traj, system, mask: str = "protein", simplified: bool = False):
    """Return DSSP-like labels for all residues."""
    return dssp(traj, system, mask=mask, simplified=simplified)[1]
