import argparse
import pickle
import time
import shutil
import os
import sys
import gzip
import logging
import json
import torch
from tqdm import tqdm
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
from geobpe.datasets import FullCathCanonicalCoordsDataset
from geobpe.tokenizer import Tokenizer
from esm.utils.structure.protein_chain import ProteinChain
from geobpe.plotting import get_codebook_utility
from concurrent.futures import ProcessPoolExecutor
from geobpe.utils import load_args_from_txt, validate_args_match, str2bool, open_pdb


def _effective_cpus() -> int:
    """Return the number of CPUs *actually* available to this task."""
    if "SLURM_CPUS_PER_TASK" in os.environ:        
        n = int(os.environ["SLURM_CPUS_PER_TASK"])
        print(f"SLURM_CPUS_PER_TASK={n}")
        return n
    try:                                   # Linux cpuset / cgroups
        return len(os.sched_getaffinity(0))
    except AttributeError:
        return os.cpu_count() or 1

def int_or_tuple(value: str):
    # normalise separators
    parts = value.replace(',', ':').split(':')
    try:
        nums = list(map(int, parts))
    except ValueError:
        raise argparse.ArgumentTypeError("expected int or int:int")
    if len(nums) == 1:
        return nums[0]
    if len(nums) == 2:
        start, end = nums
        if start < 0 or end < 0 or start >= end:
            raise argparse.ArgumentTypeError("need 0 <= start < end")
        return (start, end)

    raise argparse.ArgumentTypeError("expected at most two integers")

def add_parser(parser: argparse.ArgumentParser) -> None:
    # folder
    parser.add_argument("--src-pkl")
    parser.add_argument("--base-dir", type=str, default="./")
    parser.add_argument("--save-dir")
    parser.add_argument("--log-dir", type=str, default="logs", 
                        help="Directory where log files will be saved.")    
    parser.add_argument("--data-dir")
    # data params
    parser.add_argument(
        "--toy",
        default=0,
        type=int_or_tuple,
        help="int => take first N files; start:end or start,end => slice "
             "fnames[start:end] (end exclusive)."
    )
    parser.add_argument("--pad", default=512, type=int)
    parser.add_argument("--processed", type=str2bool, default=False)
    parser.add_argument("--debug", action='store_true')
    # induce args
    parser.add_argument("--append", action='store_true', help="Whether to append to src-pkl")


def _init_tokenize_worker(bpe, save_dir, processed_only):
    global BPE, SAVE_DIR, PROCESSED_ONLY
    BPE, SAVE_DIR, PROCESSED_ONLY = bpe, save_dir, processed_only
    


def tokenize_structure(args):
    """Build and tokenize a single structure inside the worker."""
    idx, fname = args
    struc = FullCathCanonicalCoordsDataset.single_compute_featurization(fname)
    if struc is None:
        print(f"skipping {idx}, {fname} because featurization failed")
        return False
    if (struc['angles']['psi']==struc['angles']['psi']).sum() < len(struc['angles']['psi'])-1:
        print(f"skipping {idx}, {fname} because of missing dihedrals")
        return False
    else:
        # detect if backbone parser matches ProteinChain.from_pdb            
        if len(struc['angles']) != len(ProteinChain.from_pdb(open_pdb(fname))):
            print(f"skipping {idx}, {fname} backbone parse does not match ProteinChain.from_pdb")
            return False

        if PROCESSED_ONLY and not os.path.exists(Path(__file__).parents[1] / "scripts" / Path(fname).relative_to(os.getcwd())):
            print(f"skipping {fname}, processed_only and this pdb is not processed")
            return False
    
    if os.path.exists(SAVE_DIR / f"{idx}.pkl"):
        try:
            stats, tok = pickle.load(open(SAVE_DIR / f"{idx}.pkl", "rb"))
            assert tok.fname == struc['fname']
            assert tok.n == len(ProteinChain.from_pdb(open_pdb(tok.fname)))
            return True
        except Exception as e:
            print(e, f"redo {fname}")
    else:
        print(f'start {fname}')
    tok = Tokenizer(struc)
    tok, stats = BPE.tokenize(tok)
    pickle.dump((stats, tok), open(SAVE_DIR / f"{idx}.pkl", "wb+"))
    return True


def plot_stats(all_stats, output_path, total_ticks=20):
    stats = {
        key: np.mean([d[key] for d in all_stats], axis=0).tolist()
        for key in all_stats[0]
    }
    rmsds, lddts, Ls = stats["rmsd"], stats["lddt"], stats["L"]
    fig, (ax1, ax_rmsd) = plt.subplots(1, 2, figsize=(16, 5)) # make figure + first (left) axis

    # ---------------- left panel : L vs K + BPR on right ---------------
    ax1.plot(Ls,
            marker='o',
            label="L vs Iter",
            linewidth=2)
    skip = (len(Ls)+total_ticks-1)//(total_ticks)    
    ax1.set_ylabel("L (# Motif-Tokens Per PDB)")    
    ax1.set_xlabel(f"Step")
    ax1.set_xticks(range(0, len(Ls), skip))
    ax1.tick_params(axis="y", labelcolor='tab:orange')

    lines1, labels1 = ax1.get_legend_handles_labels()
    ax1.legend(lines1, labels1, loc="best")
    ax1.set_title(f"L w/ {len(Ls)} BPE rounds")

    # -------- right panel: BB-RMSD (left-y) & LDDT (right-y) ----------
    ax_rmsd.plot(rmsds,
                 marker='s', linestyle='--', color='tab:red',
                 label="Backbone RMSD")
    ax_rmsd.set_xlabel(f"Step")
    ax_rmsd.set_xticks(range(0, len(Ls), skip))
    ax_rmsd.set_ylabel("Backbone RMSD (Å)", color='tab:red')
    ax_rmsd.tick_params(axis='y', labelcolor='tab:red')

    ax_lddt = ax_rmsd.twinx()                           # second y-axis
    ax_lddt.plot(lddts,
                 marker='o', linestyle='--', color='tab:blue',
                 label="LDDT (mean)")
    ax_lddt.set_ylabel("LDDT", color='tab:blue')
    ax_lddt.tick_params(axis='y', labelcolor='tab:blue')
    ax_rmsd.set_title("Backbone RMSD & LDDT vs K")

    # -------------------- annotate best points ------------------------
    best_rmsd_idx = np.argmin(rmsds)
    ax_rmsd.scatter(best_rmsd_idx, rmsds[best_rmsd_idx],
                    color='tab:red', zorder=5)
    ax_rmsd.annotate(f"Lowest RMSD: {rmsds[best_rmsd_idx]:.2f}",
                     xy=(best_rmsd_idx, rmsds[best_rmsd_idx]),
                     xytext=(10, 15), textcoords="offset points",
                     arrowprops=dict(arrowstyle="->", color='tab:red'),
                     color='tab:red')

    best_lddt_idx = np.argmax(lddts)
    ax_lddt.scatter(best_lddt_idx, lddts[best_lddt_idx],
                    color='tab:blue', zorder=5)
    ax_lddt.annotate(f"Highest LDDT: {lddts[best_lddt_idx]:.2f}",
                     xy=(best_lddt_idx, lddts[best_lddt_idx]),
                     xytext=(10, -15), textcoords="offset points",
                     arrowprops=dict(arrowstyle="->", color='tab:blue'),
                     color='tab:blue')

    # ---------------- combine legends for right panel -----------------
    h1, l1 = ax_rmsd.get_legend_handles_labels()
    h2, l2 = ax_lddt.get_legend_handles_labels()
    ax_rmsd.legend(h1 + h2, l1 + l2, loc='best')
    
    fig.tight_layout()
    plt.show()
    plt.savefig(output_path)
    return stats


def run(args: argparse.Namespace) -> int:
    logging.info(args)
    max_workers = _effective_cpus()
    logging.info(f"{max_workers} workers")
    if args.save_dir:
        save_dir = Path(args.save_dir)
    else:
        cur_time = time.time()
        save_dir = Path(args.base_dir) / f'./ckpts/{cur_time}'
        os.makedirs(save_dir, exist_ok=True)        
    
    print(f"save_dir: {save_dir}")
    src_file = Path(args.src_pkl)
    bpe = pickle.load(open(src_file, "rb"))
    # store metadata    
    args_path = save_dir / "args.txt"
    out_path = save_dir / src_file.name    
    if os.path.exists(args_path):
        print(f"loading args from {args_path}")
        loaded_args = load_args_from_txt(args_path)    
        validate_args_match(
            current   = args,
            loaded    = loaded_args,
            skip = ["auto", "save_dir"]
        )
    else:
        with open(args_path, "w") as f:
            for arg_name, arg_value in sorted(args.__dict__.items()):
                f.write(f"{arg_name}: {arg_value}\n")    
    arg_path = src_file.parent / "args.txt"
    try:
        shutil.copyfile(arg_path, save_dir / "orig_args.txt")
    except FileNotFoundError:
        print(arg_path, "not found")
    # induce
    fnames = FullCathCanonicalCoordsDataset._CathCanonicalAnglesDataset__get_pdb_fnames(args.data_dir)
    if args.toy and isinstance(args.toy, int): 
        fnames = fnames[:args.toy]
    elif args.toy:
        fnames = fnames[args.toy[0]: args.toy[1]]
    pargs = [(i, fname) for i, fname in enumerate(fnames)]
    N = len(pargs)
    tokenizers = []
    all_stats = []
    if not args.debug and max_workers:
        # ----- parallel tokenisation -----------------------------------------------------            
        with ProcessPoolExecutor(
            max_workers=max_workers,
            initializer=_init_tokenize_worker,
            initargs=(bpe, save_dir, args.processed)                                # only BPE is broadcast
        ) as pool:
            res = list(tqdm(pool.map(
                    tokenize_structure, pargs, chunksize=10),
                    total=N, desc="tokenizing mp"))
    else:
        global BPE, SAVE_DIR, PROCESSED_ONLY
        BPE, SAVE_DIR, PROCESSED_ONLY = bpe, save_dir, args.processed
        res = []
        for parg in tqdm(pargs, desc="tokenizing"):            
            res.append(tokenize_structure(parg))

    for i in tqdm(range(N), desc="loading done tokenizers"):
        if res[i]:
            assert os.path.exists(save_dir / f"{i}.pkl")
            stats, t = pickle.load(open(save_dir / f"{i}.pkl", "rb"))
            all_stats.append(stats)
            tokenizers.append(t)    
    print(f"loaded {len(tokenizers)} tokenizers")
    pickle.dump(tokenizers, open(save_dir / "tokenizers.pkl", "wb+"))
    input_ids = bpe.quantize(tokenizers)
    pickle.dump(input_ids, open(save_dir / "input_ids.pkl", "wb+"))
    # utility = get_codebook_utility(torch.as_tensor(sum(input_ids, [])), bpe.vocab_size)
    # json.dump(utility, open(save_dir / "utility.json", "w+"))
    try:
        stats = plot_stats(all_stats, save_dir / "stats.png")
        json.dump(stats, open(save_dir / "stats.json", "w+"))
    except:
        pass
    if args.append:
        if not isinstance(bpe.n, list):
            bpe.n = [bpe.n]
        bpe.n.append(len(tokenizers))
        bpe.tokenizers.extend(tokenizers)
    else:
        bpe.tokenizers = tokenizers
    # pickle.dump(bpe, open(out_path, "wb+"))
    print(os.path.abspath(out_path))
    return 0
