#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import sys
import pandas as pd
from pathlib import Path
from rdkit import Chem

# ============================================================
# CONFIG & AUTO-DETECTION
# ============================================================
# These match the headers generated by your previous scripts
TARGET_KEYS = ['Target', 'protein', 'protein_name']
LIGAND_KEYS = ['Warhead', 'Ligand', 'ligand', 'Ligand_Resolved']
SMILES_KEYS = ['SMILES', 'smiles']

def find_col(df_cols, keys):
    for k in keys:
        for col in df_cols:
            if k.lower() == col.lower():
                return col
    return None

# ============================================================
# MAIN SCRIPT
# ============================================================
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, required=False)
    parser.add_argument("--outdir", type=str, default="Target_Table")
    parser.add_argument("--auto", action="store_true", help="Skip interaction")
    args = parser.parse_args()

    # Determine input file
    if args.input is None:
        import glob
        csvs = glob.glob("*.csv")
        if not csvs:
            print("❌ No CSV files found.")
            sys.exit(1)
        # In auto mode, pick Resolved_SASA_Summary.csv if it exists
        if "Resolved_SASA_Summary.csv" in csvs:
            args.input = "Resolved_SASA_Summary.csv"
        else:
            args.input = csvs[0]

    print(f"📥 Loading: {args.input}")
    df = pd.read_csv(args.input)

    # Column Detection
    TARGET_COL = find_col(df.columns, TARGET_KEYS)
    LIGAND_COL = find_col(df.columns, LIGAND_KEYS)
    SMILES_COL = find_col(df.columns, SMILES_KEYS)

    # Interaction fallback if auto fails
    if not all([TARGET_COL, LIGAND_COL, SMILES_COL]) and not args.auto:
        print("📌 Available Columns:")
        for i, c in enumerate(df.columns): print(f"  {i:2d} → {c}")
        target_idx = int(input("👉 Enter index for TARGET: "))
        ligand_idx = int(input("👉 Enter index for LIGAND: "))
        smiles_idx = int(input("👉 Enter index for SMILES: "))
        TARGET_COL, LIGAND_COL, SMILES_COL = df.columns[target_idx], df.columns[ligand_idx], df.columns[smiles_idx]
    
    if not all([TARGET_COL, LIGAND_COL, SMILES_COL]):
        print(f"❌ Error: Could not find required columns. Found: T={TARGET_COL}, L={LIGAND_COL}, S={SMILES_COL}")
        sys.exit(1)

    print(f"✅ Using: Target={TARGET_COL}, Ligand={LIGAND_COL}, SMILES={SMILES_COL}")

    outdir = Path(args.outdir)
    outdir.mkdir(exist_ok=True)
    df_clean = df.dropna(subset=[TARGET_COL, LIGAND_COL, SMILES_COL])

    # 1-5) Standard Maps
    df_clean.groupby(TARGET_COL)[LIGAND_COL].unique().reset_index().to_csv(outdir / "Target_Ligand_Map.csv", index=False)
    df_clean.groupby(LIGAND_COL)[TARGET_COL].unique().reset_index().to_csv(outdir / "Ligand_Target_Map.csv", index=False)
    df_clean.groupby(LIGAND_COL)[SMILES_COL].unique().reset_index().to_csv(outdir / "Ligand_SMILES_Map.csv", index=False)
    df_clean.groupby(SMILES_COL)[LIGAND_COL].unique().reset_index().to_csv(outdir / "SMILES_Ligand_Map.csv", index=False)
    df_clean.groupby(TARGET_COL)[SMILES_COL].unique().reset_index().to_csv(outdir / "Target_SMILES_Map.csv", index=False)

    # 6) Per-Atom Mapping
    atom_rows = []
    print("🧬 Generating atom-level mapping...")
    for _, row in df_clean[[LIGAND_COL, SMILES_COL]].drop_duplicates().iterrows():
        mol = Chem.MolFromSmiles(row[SMILES_COL])
        if mol:
            for atom in mol.GetAtoms():
                atom_rows.append({"Ligand": row[LIGAND_COL], "AtomIndex": atom.GetIdx(), "AtomSymbol": atom.GetSymbol()})

    pd.DataFrame(atom_rows).to_csv(outdir / "Ligand_Atom_Map.csv", index=False)
    print(f"🎉 SUCCESS: All maps saved to {outdir}")

if __name__ == "__main__":
    main()