"""Steric parameter descriptor calculation."""

import glob
import json
from pathlib import Path
from typing import List, Optional, Union
import pandas as pd
import numpy as np
from rdkit import Chem
from rdkit.Chem import RWMol
from tqdm import tqdm

from qspoc.converter import xyz2coord

# Constants for steric calculations
Z_AXIS = np.array([0, 0, 1])
ORIGIN_2D = np.array([0, 0])


class StericHindranceCalculator:
    """Calculate steric hindrance parameters (Sterimol parameters)."""

    def __init__(self):
        """Initialize steric hindrance calculator."""
        self.L = np.nan
        self.B1 = np.nan
        self.B5 = np.nan
        self.CA = np.nan

    def load_data(
        self,
        coord: pd.DataFrame,
        atom_start_idx: int,
        atom_end_idx: Union[int, List[int]],
        mol_obj: object = None,
    ) -> None:
        """
        Load molecular data for steric calculation.

        Args:
            coord: DataFrame with atom coordinates and elements
            atom_start_idx: Index of start atom
            atom_end_idx: Index or list of indices for end atom(s)
            mol_obj: RDKit molecule object (optional, needed for automatic group detection)
        """
        # Validate inputs
        assert atom_start_idx != atom_end_idx, "Start and end atom indices should be different"
        assert atom_start_idx >= 0 and atom_start_idx < len(coord), "Atom indices out of range"

        if isinstance(atom_end_idx, int):
            assert atom_end_idx >= 0 and atom_end_idx < len(coord), "Atom indices out of range"
        elif isinstance(atom_end_idx, list):
            assert all(atom_end_idx[i] >= 0 and atom_end_idx[i] < len(coord) for i in range(len(atom_end_idx))), "Atom indices out of range"

        self.atom_start_idx = atom_start_idx
        self.atom_end_idx = atom_end_idx
        self.mol_obj = mol_obj

        # Calculate axis vector
        if isinstance(self.atom_end_idx, list):
            self.coord = coord
            self.axis_vec = self.coord.loc[self.atom_start_idx, "coord"] - self._get_centroid(self.coord.loc[self.atom_end_idx, "coord"])
        else:
            self.coord = coord.iloc[[self.atom_start_idx] + self._get_group_part(), :]
            self.axis_vec = self.coord.loc[self.atom_start_idx, "coord"] - self.coord.loc[self.atom_end_idx, "coord"]

        # Load van der Waals radii
        self.vdw_dict = self._load_vdw_radii_from_json()
        self.coord = self.coord.assign(vdw_radius=self.coord["element"].apply(self._get_vdw_radius))
        self.target_coord = self.coord[self.coord.index != self.atom_start_idx]

        # Shift coordinates to start atom origin
        base_coord = self.coord.loc[self.atom_start_idx, "coord"]
        self.target_coord.loc[:, "coord"] = self.target_coord["coord"].apply(lambda x: x - base_coord)

    def _get_centroid(self, coords: pd.Series) -> np.ndarray:
        """Calculate centroid of coordinates."""
        return np.mean(coords.values)

    def _get_group_part(self) -> List[int]:
        """
        Get atom indices of the group attached to end atom using RDKit.

        Returns:
            List of atom indices in the group
        """
        # Check if start and end atoms are bonded
        assert (
            self.mol_obj.GetBondBetweenAtoms(self.atom_start_idx, self.atom_end_idx) is not None
        ), f"Start ({self.atom_start_idx}) and end ({self.atom_end_idx}) atoms are not bonded"

        # Cut bond between start and end atoms
        rw_mol = RWMol(self.mol_obj)
        rw_mol.RemoveBond(self.atom_start_idx, self.atom_end_idx)
        frags = Chem.GetMolFrags(rw_mol)

        # Should split into two fragments
        assert len(frags) == 2, "Molecule should be cut into two fragments"

        # Return the fragment containing the end atom
        for frag in frags:
            if self.atom_end_idx in frag:
                return list(frag)

        raise Exception("Failed to cut molecule into two fragments")

    def _load_vdw_radii_from_json(self, vdw_filename: str = "vdw_radii.json") -> dict:
        """
        Load van der Waals radii from JSON file.

        Args:
            vdw_filename: Name of JSON file containing radii data

        Returns:
            Dictionary mapping element symbols to radii
        """
        json_file = Path(__file__).parent / vdw_filename
        with open(json_file, "r", encoding="utf-8") as f:
            vdw_data = json.load(f)

            # Build {element symbol: radius} dictionary
            vdw_dict = {}
            for item in vdw_data:
                symbol = item.get("symbol")
                radius = item.get("radius")
                if symbol and radius is not None:
                    vdw_dict[symbol] = float(radius)

        return vdw_dict

    def _get_vdw_radius(self, atom_name: str) -> float:
        """
        Get van der Waals radius for an element.

        Args:
            atom_name: Element symbol

        Returns:
            Van der Waals radius in Angstroms
        """
        radius = self.vdw_dict.get(atom_name)
        if radius is None:
            print(
                f"Unknown atomic symbol '{atom_name}', corresponding van der Waals radius not found,\
                default value 2.0A will be used"
            )
            return 2.0
        return radius

    def _angle_between(self, v1: np.ndarray, v2: np.ndarray) -> float:
        """
        Calculate angle between two vectors (in radians).

        Args:
            v1: First vector
            v2: Second vector

        Returns:
            Angle in radians
        """
        dot = np.dot(v1, v2)
        norm = np.linalg.norm(v1) * np.linalg.norm(v2)
        if norm == 0:
            return 0.0
        return np.arccos(np.clip(dot / norm, -1.0, 1.0))

    def _calc_adj_opp(self, c: np.ndarray) -> tuple[float, float]:
        """
        Calculate parallel (adj) and perpendicular (opp) distances of atom to Sterimol axis.

        Args:
            c: Atom coordinates

        Returns:
            Tuple of (adj, opp) distances
        """
        angle = self._angle_between(self.axis_vec, c)
        hypotenuse = np.linalg.norm(c)
        adj = hypotenuse * np.cos(angle)  # distance along axial direction
        opp = hypotenuse * np.sin(angle)  # distance perpendicular to axial direction
        return adj, opp

    def _rotate_to_z_axis(self, cart: np.ndarray) -> np.ndarray:
        """
        Rotate coordinates to align Sterimol axis with z-axis (for B1 calculation).

        Args:
            cart: Cartesian coordinates

        Returns:
            Rotated coordinates
        """
        # Calculate rotation matrix based on original algorithm's rotrel function
        ax = np.cross(self.axis_vec, Z_AXIS)
        ang = self._angle_between(self.axis_vec, Z_AXIS)

        if np.linalg.norm(ax) == 0:
            return cart

        ax_norm = ax / np.linalg.norm(ax)
        ux, uy, uz = ax_norm
        cos_ang = np.cos(ang)
        sin_ang = np.sin(ang)

        # Rotation matrix
        rot_mat = np.array(
            [
                [cos_ang + ux**2 * (1 - cos_ang), ux * uy * (1 - cos_ang) - uz * sin_ang, ux * uz * (1 - cos_ang) + uy * sin_ang],
                [uy * ux * (1 - cos_ang) + uz * sin_ang, cos_ang + uy**2 * (1 - cos_ang), uy * uz * (1 - cos_ang) - ux * sin_ang],
                [uz * ux * (1 - cos_ang) - uy * sin_ang, uz * uy * (1 - cos_ang) + ux * sin_ang, cos_ang + uz**2 * (1 - cos_ang)],
            ]
        )
        return np.dot(rot_mat, cart)

    def _get_circle_tangents(self, c_atom1: np.ndarray, c_atom2: np.ndarray, r1: float, r2: float) -> List[List[np.ndarray]]:
        """
        Get external tangents between two circles in 2D.

        Args:
            c_atom1: Center of first circle
            c_atom2: Center of second circle
            r1: Radius of first circle
            r2: Radius of second circle

        Returns:
            List of tangent segments [[p1a, p2a], [p1b, p2b]]
        """
        dc = c_atom2 - c_atom1  # Vector between circle centers
        d_sq = dc @ dc  # Squared distance (avoid sqrt for precision)
        d = np.sqrt(d_sq)  # Distance between centers

        # Handle concentric circles (allow small floating point deviation)
        if d < 1e-6:
            return []

        u = dc / d  # Unit vector along line joining circle centers
        v = np.array([-u[1], u[0]])
        k = r1 - r2
        h_sq = d_sq - k**2

        if h_sq < -1e-9:
            external_tangents = []
        else:
            h = np.sqrt(max(h_sq, 0.0))
            u1a = (k * u + h * v) / d
            u1b = (k * u - h * v) / d
            p1a = c_atom1 + r1 * u1a
            p2a = c_atom2 + r2 * u1a
            p1b = c_atom1 + r1 * u1b
            p2b = c_atom2 + r2 * u1b
            external_tangents = [[p1a, p2a], [p1b, p2b]]

        return external_tangents

    def _dist_2d(self, p: np.ndarray, v_start: np.ndarray, v_end: np.ndarray) -> float:
        """
        Calculate distance from point to line segment in 2D.

        Args:
            p: Point coordinates
            v_start: Start of line segment
            v_end: End of line segment

        Returns:
            Perpendicular distance from point to line segment
        """
        vect1 = np.subtract(p, v_start)
        vect2 = np.subtract(v_start, v_end)
        ang = np.arccos(vect1 @ vect2 / np.linalg.norm(vect1) / np.linalg.norm(vect2))
        return np.sin(ang) * np.linalg.norm(vect1)

    def _calc_B1(self) -> float:
        """
        Calculate minimum lateral width B1 (shortest tangent between atoms in 2D plane).

        Returns:
            Minimum lateral width B1 in Angstroms
        """
        # Convert coordinates to 2D plane with start atom as origin and axis as z-axis
        z_origin = self.coord.loc[self.atom_start_idx, "coord"]
        carts_2d = self.target_coord.copy()
        carts_2d["coord"] = carts_2d["coord"].apply(lambda x: self._rotate_to_z_axis(x - z_origin))
        carts_2d["coord"] = carts_2d["coord"].apply(lambda x: x[:2])

        # Find valid tangent vectors
        valid_vecs = {"vec": [], "dist": []}
        for atom1_idx in carts_2d.index:
            for atom2_idx in carts_2d.index:
                if atom1_idx > atom2_idx:
                    continue

                carts_2d_atom1 = carts_2d.loc[atom1_idx, "coord"]
                carts_2d_atom2 = carts_2d.loc[atom2_idx, "coord"]
                radii_atom1, radii_atom2 = (carts_2d.loc[atom1_idx, "vdw_radius"], carts_2d.loc[atom2_idx, "vdw_radius"])

                vecs = self._get_circle_tangents(carts_2d_atom1, carts_2d_atom2, radii_atom1, radii_atom2)

                for vec in vecs:
                    is_valid_vec = True
                    for atomk_idx in carts_2d.index:
                        if atomk_idx == atom1_idx or atomk_idx == atom2_idx:
                            continue
                        dist = self._dist_2d(carts_2d.loc[atomk_idx, "coord"], vec[0], vec[1])
                        if dist <= carts_2d.loc[atomk_idx, "vdw_radius"] - 1e-6:
                            # Tangent passes through an atom
                            is_valid_vec = False
                            break

                    if is_valid_vec:
                        valid_vecs["vec"].append(vec)
                        valid_vecs["dist"].append(self._dist_2d(ORIGIN_2D, vec[0], vec[1]))

        # Plot if needed (commented out for production)
        # plot_all_atom_and_vecs(carts_2d, valid_vecs["vec"])

        if valid_vecs["vec"]:
            return min(valid_vecs["dist"])
        else:
            return carts_2d["vdw_radius"].max()

    def _calc_V_bur(self, num_samples: int = 100000, sample_radius: float = 3.5) -> float:
        """
        Calculate %V_bur (percent buried volume).

        Args:
            num_samples: Number of Monte Carlo samples
            sample_radius: Radius of sampling sphere

        Returns:
            Percent buried volume (0-100)
        """
        centers = np.vstack(self.target_coord["coord"].values)
        radii = self.target_coord["vdw_radius"].values

        # Generate random points in sphere
        norm_points = np.random.normal(size=(num_samples, 3))
        norms = np.linalg.norm(norm_points, axis=1, keepdims=True)
        norms[norms == 0] = 1e-10
        unit_sphere_points = norm_points / norms
        s = np.random.uniform(0, 1, size=(num_samples, 1))
        scaling_factors = sample_radius * (s ** (1 / 3))
        points = unit_sphere_points * scaling_factors

        # Check which points are inside any atom's van der Waals sphere
        dist_sq = np.sum((points[:, np.newaxis, :] - centers) ** 2, axis=2)
        in_any_ball = np.any(dist_sq <= radii**2, axis=1)
        num_hits = np.sum(in_any_ball)

        return float(num_hits / num_samples)

    def _calculate_parameters(self) -> None:
        """Execute all parameter calculations for a single axis."""
        # Calculate L (axial length) and B5 (maximum lateral width)
        adj_list = []
        opp_list = []
        for _, c in self.target_coord.iterrows():
            adj, opp = self._calc_adj_opp(c["coord"])
            adj_list.append(np.abs(adj) + c["vdw_radius"])  # Parallel distance + vdW
            opp_list.append(np.abs(opp) + c["vdw_radius"])  # Perpendicular distance + vdW

        self.L = max(adj_list) + 0.4 if adj_list else np.nan  # 0.4Å correction from original algorithm
        self.B5 = max(opp_list) if opp_list else np.nan
        self.B1 = self._calc_B1()
        self.V_Bur = self._calc_V_bur()

    def _get_results(self, idx: Optional[int] = None) -> dict:
        """
        Return formatted results.

        Args:
            idx: Optional index to append to parameter names

        Returns:
            Dictionary with steric parameters
        """
        # Round results
        self.L = round(self.L, 3) if not np.isnan(self.L) else self.L
        self.B5 = round(self.B5, 3) if not np.isnan(self.B5) else self.B5
        self.B1 = round(self.B1, 3) if not np.isnan(self.B1) else self.B1
        self.V_Bur = round(self.V_Bur, 3) if not np.isnan(self.V_Bur) else self.V_Bur

        if idx is not None:
            return {f"L-{idx}": self.L, f"B1-{idx}": self.B1, f"B5-{idx}": self.B5, f"%V_Bur-{idx}": self.V_Bur}
        else:
            return {"L": self.L, "B1": self.B1, "B5": self.B5, "%V_Bur": self.V_Bur}

    def calculate(self, calc_data_df: pd.DataFrame, temp_coords_dir: Path) -> pd.DataFrame:
        """
        Calculate steric parameters for all molecules.

        Args:
            calc_data_df: DataFrame containing calculation data with columns:
                - smiles_list: SMILES strings
                - name_list: Unique identifiers for molecules
                - axis_lists: List of axis definitions for each molecule
            temp_coords_dir: Directory containing optimized coordinate files

        Returns:
            DataFrame with steric parameters
        """
        n_sp = len(calc_data_df.loc[0, "axis_lists"])
        sp_columns = [[f"L-{idx}", f"B1-{idx}", f"B5-{idx}", f"%V_Bur-{idx}"] for idx in range(n_sp)]
        sp_columns = [item for sublist in sp_columns for item in sublist]
        sp_df = pd.DataFrame(columns=sp_columns)

        pbar = tqdm(
            zip(calc_data_df["smiles_list"], calc_data_df["name_list"], calc_data_df["axis_lists"]),
            total=len(calc_data_df),
            desc="[Steric] Parameters",
            ncols=100,
            unit="mol",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
        )

        for s, tag, axis in pbar:
            pbar.set_postfix_str(f"Processing {tag}")
            escaped_tag = glob.escape(tag)
            opt_file = [f for f in Path(temp_coords_dir).glob(f"{escaped_tag}-opt_with_*.xyz")]
            opt_method = [of.stem.split("_")[-1] for of in opt_file]
            if len(opt_method) > 1:
                opt_method = "g16" if "g16" in opt_method else "xtb"
            else:
                opt_method = opt_method[0]
            coord = xyz2coord(temp_coords_dir / f"{tag}-opt_with_{opt_method}.xyz")

            mol_obj = Chem.MolFromSmiles(s)
            steric_param_info = {}
            for i, ax in enumerate(axis):
                calculator = StericHindranceCalculator()
                calculator.load_data(coord=coord, atom_start_idx=ax[0], atom_end_idx=ax[1], mol_obj=mol_obj)
                calculator._calculate_parameters()
                steric_param_info.update(calculator._get_results(i))
            sp_df.loc[tag, :] = pd.Series(steric_param_info)

        return sp_df
