#!/usr/bin/env python
# coding: utf-8
"""
Compute ESMFold structure statistics from PDB and pre-generated DSSP files.

PDB folder (--input_files): sequence, GRAVY, pI, ASA, Rg, NtoCdistance, pLDDT.
DSSP folder (--dssp_dir): general + per-type secondary structure (raw DSSP codes, no angle correction).

Example:
  python esmfold_processing_DSSP.py \\
    --input_files /path/to/pdbs \\
    --dssp_dir /path/to/dssp

Output CSV columns (sequence_id = PDB/DSSP filename stem):
  sequence_id, sequence, GRAVY, pI,
  helix, strand, disorder, structured,
  alpha-helix, helix-3, helix-5, helix-PPII, betabridge, turn, bend, loops,
  ASA, Rg, NtoCdistance, pLDDT
"""

OUTPUT_COLUMNS = [
    "sequence_id",
    "sequence",
    "GRAVY",
    "pI",
    "helix",
    "strand",
    "disorder",
    "structured",
    "alpha-helix",
    "helix-3",
    "helix-5",
    "helix-PPII",
    "betabridge",
    "turn",
    "bend",
    "loops",
    "ASA",
    "Rg",
    "NtoCdistance",
    "pLDDT",
]

# DSSP one-letter codes per category (direct assignment, no phi/psi correction)
SS_CATEGORIES = {
    "helix": {"G", "H", "I", "P"},
    "strand": {"E"},
    "disordered": {"-", "C", " ", "B", "T", "S"},
    "structured": {"G", "H", "I", "E", "P"},
    "alpha_helix": {"H"},
    "helix-3": {"G"},
    "helix-5": {"I"},
    "helix-PPII": {"P"},
    "betabridge": {"B"},
    "turn": {"T"},
    "bend": {"S"},
    "loops": {"B", "T", "S"},
}

import argparse
import glob
import math
import os
from statistics import mean

import numpy as np
import pandas as pd
from Bio.PDB import PDBParser
from Bio.PDB.SASA import ShrakeRupley
from Bio.SeqUtils import seq1
from Bio.SeqUtils.IsoelectricPoint import IsoelectricPoint

# Kyte-Doolittle hydropathy index
HYDROPATHY = {
    "A": 1.8,
    "C": 2.5,
    "D": -3.5,
    "E": -3.5,
    "F": 2.8,
    "G": -0.4,
    "H": -3.2,
    "I": 4.5,
    "K": -3.9,
    "L": 3.8,
    "M": 1.9,
    "N": -3.5,
    "P": -1.6,
    "Q": -3.5,
    "R": -4.5,
    "S": -0.8,
    "T": -0.7,
    "V": 4.2,
    "W": -0.9,
    "Y": -1.3,
}

try:
    from Bio.PDB.DSSP import dssp_dict_from_dssp_file
except ImportError:
    dssp_dict_from_dssp_file = None


def list_structure_files(directory, extensions):
    """Return sorted file paths for given extensions in a single folder."""
    paths = []
    for ext in extensions:
        paths.extend(glob.glob(os.path.join(directory, f"*{ext}")))
    return sorted(paths)


def file_stem(filepath):
    return os.path.splitext(os.path.basename(filepath))[0]


def read_pdbs(pdb_dir):
    return list_structure_files(pdb_dir, [".pdb"])


def pair_pdb_with_dssp(pdbs, dssp_dir):
    """
    Match each PDB to a DSSP file in dssp_dir by basename (variant.pdb -> variant.dssp).
    Returns paired lists in the same order and lists of unmatched names.
    """
    paired_pdbs = []
    paired_dssp = []
    missing_dssp = []

    for pdb_path in pdbs:
        stem = file_stem(pdb_path)
        dssp_path = os.path.join(dssp_dir, f"{stem}.dssp")
        if os.path.isfile(dssp_path):
            paired_pdbs.append(pdb_path)
            paired_dssp.append(dssp_path)
        else:
            missing_dssp.append(stem)

    pdb_stems = {file_stem(p) for p in pdbs}
    orphan_dssp = [
        file_stem(p)
        for p in list_structure_files(dssp_dir, [".dssp"])
        if file_stem(p) not in pdb_stems
    ]

    return paired_pdbs, paired_dssp, missing_dssp, orphan_dssp


def standard_amino_acids(sequence):
    return "".join(aa for aa in str(sequence).upper() if aa in HYDROPATHY)


def compute_gravy(sequence):
    """Mean Kyte-Doolittle GRAVY score for a protein sequence."""
    residues = [HYDROPATHY[aa] for aa in standard_amino_acids(sequence)]
    if not residues:
        return np.nan
    return float(np.mean(residues))


def compute_pI(sequence):
    """Isoelectric point (pI) from the protein sequence."""
    clean = standard_amino_acids(sequence)
    if not clean:
        return np.nan
    try:
        return round(IsoelectricPoint(clean).pi(), 3)
    except Exception:
        return np.nan


def sequence_from_pdb(pdb_path, chain_id="A"):
    """One-letter amino-acid sequence from the first model (chain A by default)."""
    parser = PDBParser(QUIET=1)
    structure = parser.get_structure(file_stem(pdb_path), pdb_path)
    model = structure[0]
    if chain_id in model:
        chain = model[chain_id]
    else:
        chain = next(iter(model))

    letters = []
    for residue in chain:
        if residue.id[0] != " ":
            continue
        try:
            letters.append(seq1(residue.get_resname()))
        except KeyError:
            continue
    return "".join(letters)


def sequence_metrics_bulk(entries_pdb):
    """Extract sequence from each PDB and compute GRAVY and isoelectric point."""
    rows = {}
    for entry in entries_pdb:
        prot_id = file_stem(entry)
        sequence = sequence_from_pdb(entry)
        rows[prot_id] = {
            "sequence": sequence,
            "GRAVY": compute_gravy(sequence),
            "pI": compute_pI(sequence),
        }
    df = pd.DataFrame.from_dict(rows, orient="index")
    df.index.name = "sequence_id"
    return df


def Rg(filename):
    """Radius of gyration (Å) from a PDB file."""
    coord = []
    mass = []
    with open(filename) as structure_file:
        for line in structure_file:
            try:
                parts = line.split()
                x = float(parts[6])
                y = float(parts[7])
                z = float(parts[8])
                coord.append([x, y, z])
                element = parts[-1]
                if element == "C":
                    mass.append(12.0107)
                elif element == "O":
                    mass.append(15.9994)
                elif element == "N":
                    mass.append(14.0067)
                elif element == "S":
                    mass.append(32.065)
            except (IndexError, ValueError):
                pass
    xm = [(m * i, m * j, m * k) for (i, j, k), m in zip(coord, mass)]
    tmass = sum(mass)
    rr = sum(
        mi * i + mj * j + mk * k
        for (i, j, k), (mi, mj, mk) in zip(coord, xm)
    )
    mm = sum((sum(i) / tmass) ** 2 for i in zip(*xm))
    return round(math.sqrt(rr / tmass - mm), 3)


def calc_residue_dist(residue_one, residue_two):
    """C-alpha distance between two residues."""
    diff_vector = residue_one["CA"].coord - residue_two["CA"].coord
    return np.sqrt(np.sum(diff_vector * diff_vector))


def NtoC_distance(entries_pdb, chain_id="A"):
    """C-alpha distance between N- and C-terminal residues (Å)."""
    parser = PDBParser(QUIET=1)
    distances = {}
    for entry in entries_pdb:
        prot_id = file_stem(entry)
        structure = parser.get_structure(prot_id, entry)
        model = structure[0]
        chain = model[chain_id] if chain_id in model else next(iter(model))
        residues = [r for r in chain if r.id[0] == " " and "CA" in r]
        if len(residues) < 2:
            distances[prot_id] = np.nan
            continue
        distances[prot_id] = round(
            calc_residue_dist(residues[0], residues[-1]), 3
        )
    return pd.DataFrame.from_dict(
        distances, orient="index", columns=["NtoCdistance"]
    )


def rg_bulk(entries_pdb):
    rgs = {file_stem(entry): Rg(entry) for entry in entries_pdb}
    return pd.DataFrame.from_dict(rgs, orient="index", columns=["Rg"])


def get_plddt(entries_pdb):
    parser = PDBParser()
    plddts_mean = {}
    for entry in entries_pdb:
        prot_id = file_stem(entry)
        structure = parser.get_structure(entry, entry)
        plddts = [a.get_bfactor() for a in structure.get_atoms()]
        plddts_mean[prot_id] = mean(plddts)
    return pd.DataFrame.from_dict(plddts_mean, orient="index", columns=["pLDDT"])


def compute_asa(entries_pdb):
    sr = ShrakeRupley()
    asa_models = {}
    for entry in entries_pdb:
        prot_id = file_stem(entry)
        structure = PDBParser().get_structure(entry, entry)
        sr.compute(structure, level="S")
        asa_models[prot_id] = round(structure.sasa, 2)
    return pd.DataFrame.from_dict(asa_models, orient="index", columns=["ASA"])


def _parse_dssp_manual(filepath):
    """Parse classic DSSP output (e.g. mkdssp --output-format=dssp)."""
    ss_dict = {}
    with open(filepath) as handle:
        for line in handle:
            if len(line) < 17:
                continue
            if line.startswith("#"):
                continue
            aa = line[13]
            ss = line[16]
            if aa == "!":
                continue
            if not aa.isalpha():
                continue
            try:
                res_id = int(line[5:10].strip())
            except ValueError:
                res_id = len(ss_dict) + 1
            ss_dict[res_id] = [ss]
    return ss_dict


def parse_dssp(filepath):
    """
    Per-residue DSSP dict: residue_key -> [secondary_structure_code].
    Uses BioPython when available (SS = tuple index 2); otherwise parses the .dssp file.
    """
    if dssp_dict_from_dssp_file is not None:
        try:
            dssp_dict = dssp_dict_from_dssp_file(filepath)[0]
            return {key: [value[2]] for key, value in dssp_dict.items()}
        except Exception:
            pass
    return _parse_dssp_manual(filepath)


def _ss_code(value):
    return value[0]


def get_SS_proportion(ss_dict, category):
    """Fraction of residues in a secondary-structure category (raw DSSP codes)."""
    codes = SS_CATEGORIES.get(category)
    if codes is None:
        raise ValueError(f"Unknown SS category: {category}")
    if not ss_dict:
        return 0.0
    count = sum(1 for value in ss_dict.values() if _ss_code(value) in codes)
    return round(count / len(ss_dict), 2)


def entries_proportions(dssp_entries):
    dic = {}
    for entry in dssp_entries:
        prot_id = file_stem(entry)
        entry_ss = parse_dssp(entry)
        dic[prot_id] = {
            "helix": get_SS_proportion(entry_ss, "helix"),
            "strand": get_SS_proportion(entry_ss, "strand"),
            "disorder": get_SS_proportion(entry_ss, "disordered"),
            "structured": get_SS_proportion(entry_ss, "structured"),
            "alpha-helix": get_SS_proportion(entry_ss, "alpha_helix"),
            "helix-3": get_SS_proportion(entry_ss, "helix-3"),
            "helix-5": get_SS_proportion(entry_ss, "helix-5"),
            "helix-PPII": get_SS_proportion(entry_ss, "helix-PPII"),
            "betabridge": get_SS_proportion(entry_ss, "betabridge"),
            "turn": get_SS_proportion(entry_ss, "turn"),
            "bend": get_SS_proportion(entry_ss, "bend"),
            "loops": get_SS_proportion(entry_ss, "loops"),
        }
    df = pd.DataFrame.from_dict(dic, orient="index")
    df.index.name = "sequence_id"
    return df


def main():
    parser = argparse.ArgumentParser(
        description="ESMFold stats from PDB files and pre-generated DSSP files."
    )
    parser.add_argument(
        "--input_files",
        type=str,
        required=True,
        help="Folder containing .pdb files.",
    )
    parser.add_argument(
        "--dssp_dir",
        type=str,
        required=True,
        help="Folder with pre-generated .dssp files (matched to PDBs by filename stem).",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output CSV path (default: stats.csv in cwd).",
    )
    args = parser.parse_args()

    pdb_dir = os.path.abspath(args.input_files)
    dssp_dir = os.path.abspath(args.dssp_dir)

    pdbs = read_pdbs(pdb_dir)
    if not pdbs:
        raise SystemExit(f"No .pdb files found in {pdb_dir}")

    pdbs, dssp_files, missing_dssp, orphan_dssp = pair_pdb_with_dssp(pdbs, dssp_dir)

    if missing_dssp:
        print(
            f"Warning: no matching .dssp in {dssp_dir} for "
            f"{len(missing_dssp)} PDB(s), skipping: {', '.join(missing_dssp)}"
        )
    if orphan_dssp:
        print(
            f"Warning: .dssp without matching .pdb in {pdb_dir}, ignoring: "
            f"{', '.join(orphan_dssp)}"
        )
    if not pdbs:
        raise SystemExit(
            f"No PDB/DSSP pairs found. PDBs in {pdb_dir}, DSSP in {dssp_dir} "
            "(expected <stem>.pdb and <stem>.dssp)."
        )

    print(f"Processing {len(pdbs)} structure(s) (PDB: {pdb_dir}, DSSP: {dssp_dir})")

    ss_df = entries_proportions(dssp_files)
    seq_df = sequence_metrics_bulk(pdbs)
    asa_df = compute_asa(pdbs)
    rgs_df = rg_bulk(pdbs)
    ntoc_df = NtoC_distance(pdbs)
    plddts_df = get_plddt(pdbs)

    for df in (asa_df, rgs_df, ntoc_df, plddts_df):
        df.index.name = "sequence_id"

    all_stats = (
        ss_df.join(seq_df, how="inner")
        .join(asa_df, how="inner")
        .join(rgs_df, how="inner")
        .join(ntoc_df, how="inner")
        .join(plddts_df, how="inner")
        .reset_index()
    )

    n_missing_gravy = all_stats["GRAVY"].isna().sum()
    if n_missing_gravy:
        print(f"Warning: {n_missing_gravy} structure(s) have no GRAVY (empty sequence).")
    n_missing_pi = all_stats["pI"].isna().sum()
    if n_missing_pi:
        print(f"Warning: {n_missing_pi} structure(s) have no pI (empty sequence).")
    n_missing_ntoc = all_stats["NtoCdistance"].isna().sum()
    if n_missing_ntoc:
        print(
            f"Warning: {n_missing_ntoc} structure(s) have no NtoCdistance "
            "(fewer than 2 C-alpha atoms)."
        )

    all_stats = all_stats[OUTPUT_COLUMNS]

    print("Output columns:", ", ".join(OUTPUT_COLUMNS))
    print(all_stats)

    out_path = args.output if args.output else "stats.csv"

    all_stats.to_csv(out_path, index=False)
    print(f"Wrote {out_path}")


if __name__ == "__main__":
    main()