"""Local BERT fine tuning, validation checkpoint selection and held-out evaluation."""
import argparse
import copy
import json
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import pandas as pd
import torch
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix, roc_auc_score, f1_score
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup
from common.io import ROOT, dump, seed_all, environment, sha256


def output_directories(args):
    """Keep the original seed-42 defaults; isolate every optional repeat run."""
    original_results = (ROOT / "results/fine_tuning").resolve()
    original_model = (ROOT / "models/disaster-bert").resolve()
    custom = args.result_dir is not None or args.model_output_dir is not None
    base = Path(args.base_model_dir)
    base = (base if base.is_absolute() else ROOT / base).resolve()
    original_base = base == (ROOT / "models/bert-mini").resolve()
    original_id = args.model_id == "google/bert_uncased_L-4_H-256_A-4"
    if not custom and args.seed == 42 and original_base and original_id:
        return original_results, original_model
    if args.result_dir is None or args.model_output_dir is None:
        raise ValueError("Repeat/custom runs require both --result-dir and --model-output-dir")

    def resolve(value):
        path = Path(value)
        return (path if path.is_absolute() else ROOT / path).resolve()

    results, model = resolve(args.result_dir), resolve(args.model_output_dir)
    for target in (results, model):
        for protected in (original_results, original_model):
            if target == protected or protected in target.parents or target in protected.parents:
                raise ValueError(f"Custom output overlaps the protected original artifacts: {target}")
        if target.exists() and (not target.is_dir() or any(target.iterdir())):
            raise ValueError(f"Custom output must be an empty/new directory: {target}")
    if results == model or results in model.parents or model in results.parents:
        raise ValueError("Result and model output directories must be separate")
    return results, model

def metrics(y, pred, prob=None):
    p, r, f, _ = precision_recall_fscore_support(y, pred, average="binary", zero_division=0)
    m = dict(accuracy=float(accuracy_score(y, pred)), precision=float(p), recall=float(r), f1=float(f),
             macro_f1=float(f1_score(y, pred, average="macro")), confusion_matrix=confusion_matrix(y, pred, labels=[0,1]).tolist())
    if prob is not None:
        m["roc_auc"] = float(roc_auc_score(y, prob))
    return m

def bootstrap_ci(y, pred, seed=42):
    rng = np.random.default_rng(seed)
    scores = []
    y, pred = np.asarray(y), np.asarray(pred)
    for _ in range(1000):
        ix = rng.integers(0, len(y), len(y))
        scores.append(f1_score(y[ix], pred[ix], zero_division=0))
    return np.quantile(scores, [0.025, 0.975]).tolist()

class Tweets(Dataset):
    def __init__(self, df, tokenizer, max_length):
        self.enc = tokenizer(df.text.tolist(), truncation=True, max_length=max_length, padding="max_length")
        self.labels = df.target.tolist()
    def __len__(self):
        return len(self.labels)
    def __getitem__(self, i):
        return {**{k: torch.tensor(v[i]) for k,v in self.enc.items()}, "labels": torch.tensor(self.labels[i])}

@torch.inference_mode()
def predict(model, loader, device):
    model.eval()
    probs = []
    for batch in loader:
        batch = {k:v.to(device) for k,v in batch.items()}
        probs.extend(model(**batch).logits.softmax(-1)[:,1].cpu().tolist())
    return np.asarray(probs)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=5)
    ap.add_argument("--batch-size", type=int, default=32)
    ap.add_argument("--lr", type=float, default=3e-5)
    ap.add_argument("--max-length", type=int, default=96)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--result-dir", help="Optional fresh result directory; relative paths use the assignment root")
    ap.add_argument("--model-output-dir", help="Optional fresh checkpoint directory; relative paths use the assignment root")
    ap.add_argument("--base-model-dir", default="models/bert-mini", help="Local base model/tokenizer directory; relative paths use the assignment root")
    ap.add_argument("--model-id", default="google/bert_uncased_L-4_H-256_A-4", help="Base model identity recorded in the results")
    args = ap.parse_args()
    try:
        out, model_out = output_directories(args)
    except ValueError as error:
        ap.error(str(error))
    seed_all(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    out.mkdir(parents=True, exist_ok=True)
    split_hashes = {name: sha256(ROOT / f"data/splits/{name}.csv") for name in ["train", "validation", "test"]}
    parts = {name: pd.read_csv(ROOT / f"data/splits/{name}.csv", keep_default_na=False) for name in ["train","validation","test"]}
    train, val, test = (parts[n] for n in ["train","validation","test"])
    base_model = Path(args.base_model_dir)
    base_model = base_model if base_model.is_absolute() else ROOT / base_model
    tok = AutoTokenizer.from_pretrained(base_model, local_files_only=True)
    loaders = {n:DataLoader(Tweets(df,tok,args.max_length), batch_size=args.batch_size, shuffle=n=="train", num_workers=0) for n,df in parts.items()}
    model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=2,
        id2label={0:"Not disaster",1:"Disaster"}, label2id={"Not disaster":0,"Disaster":1}, local_files_only=True).to(device)
    optimizer = torch.optim.AdamW(model.parameters(),lr=args.lr,weight_decay=0.01)
    steps = args.epochs * len(loaders["train"])
    scheduler = get_linear_schedule_with_warmup(optimizer, int(0.1*steps), steps)
    history, best_score, best_state, best_epoch = [], -1, None, None
    start = time.perf_counter()
    for epoch in range(1,args.epochs+1):
        model.train()
        total_loss = 0
        for batch in loaders["train"]:
            batch = {k:v.to(device) for k,v in batch.items()}
            optimizer.zero_grad(set_to_none=True)
            loss = model(**batch).loss
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)
            optimizer.step(); scheduler.step()
            total_loss += loss.item()*len(batch["labels"])
        prob = predict(model,loaders["validation"],device)
        result = dict(epoch=epoch, train_loss=total_loss/len(train), validation=metrics(val.target,prob>=0.5,prob))
        history.append(result)
        dump(out / "history.json",history)
        print(json.dumps(result),flush=True)
        if result["validation"]["f1"] > best_score:
            best_score=result["validation"]["f1"]; best_epoch=epoch
            best_state={k:v.detach().cpu().clone() for k,v in model.state_dict().items()}
    model.load_state_dict(best_state)
    model.save_pretrained(model_out)
    tok.save_pretrained(model_out)
    prob=predict(model,loaders["test"],device)
    pred=(prob>=0.5).astype(int)
    test_results=metrics(test.target,pred,prob)
    test_results["bootstrap_f1_ci95"]=bootstrap_ci(test.target,pred)
    baselines={}
    majority=np.full(len(test),int(train.target.mode()[0]))
    baselines["majority"]=metrics(test.target,majority)
    pipeline=make_pipeline(TfidfVectorizer(lowercase=True,ngram_range=(1,2),min_df=2,sublinear_tf=True),
                          LogisticRegression(C=2.0,max_iter=1000,random_state=args.seed))
    pipeline.fit(train.text,train.target)
    base_prob=pipeline.predict_proba(test.text)[:,1]
    base_pred=(base_prob>=0.5).astype(int)
    baselines["tfidf_logistic"]=metrics(test.target,base_pred,base_prob)
    result=dict(config=vars(args),device=device,environment=environment(),model=args.model_id,
        data_split_sha256=split_hashes, model_output_dir=str(model_out), base_model_dir=str(base_model.resolve()),
        trainable_parameters=sum(p.numel() for p in model.parameters()),best_epoch=best_epoch,
        best_validation_f1=best_score,test=test_results,baselines=baselines,seconds=time.perf_counter()-start,
        selection="best validation F1; fixed threshold 0.5; test evaluated after selection")
    dump(out / "metrics.json",result)
    pd.DataFrame(dict(id=test.id,text=test.text,gold=test.target,prediction=pred,prob_disaster=prob,
                      tfidf_prediction=base_pred,tfidf_prob=base_prob)).to_csv(out / "test_predictions.csv",index=False,encoding="utf-8-sig")
    errors=test.assign(prediction=pred,prob_disaster=prob)
    errors=errors[errors.target!=errors.prediction]
    errors["confidence"]=np.maximum(errors.prob_disaster,1-errors.prob_disaster)
    errors.sort_values("confidence",ascending=False).to_csv(out / "error_cases.csv",index=False,encoding="utf-8-sig")
    print(json.dumps(result,ensure_ascii=False,indent=2),flush=True)

if __name__ == "__main__":
    main()