import os
# Force PyTorch to ignore CUDA completely so 0 MB GPU VRAM is allocated
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import sys
import json
import shutil
import subprocess
import torch
import torch.nn as nn
import numpy as np
from collections import Counter
from sklearn.model_selection import StratifiedGroupKFold
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    DataCollatorWithPadding,
    TrainerCallback
)

# Label Mapping for Vox v7 Gate 3 Operational Edge Ontology
LABEL2ID = {
    "SHAPES": 0,
    "DEPENDS_ON": 1,
    "CONFLICTS_WITH": 2,
    "NONE": 3
}
ID2LABEL = {v: k for k, v in LABEL2ID.items()}

DATASET_PATH = os.path.expanduser("~/.vox/sandbox/dataset.json")
OUTPUT_DIR = os.path.expanduser("~/.vox/sandbox/output_6e")
METRICS_LOG_PATH = os.path.expanduser("~/.vox/sandbox/training_metrics_6e.json")
ONNX_EXPORT_PATH = os.path.expanduser("~/.vox/sandbox/model_quantized_v2.onnx")
MODEL_NAME = "answerdotai/ModernBERT-base"

def format_input(fact_a, fact_b, context):
    return f"Fact A: {fact_a} | Fact B: {fact_b} | Context: {context}"

class WeightedLossTrainer(Trainer):
    def __init__(self, class_weights=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if class_weights is not None:
            self.class_weights = torch.tensor(class_weights, dtype=torch.float32).to(self.args.device)
        else:
            self.class_weights = None

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        if self.class_weights is not None:
            loss_fct = nn.CrossEntropyLoss(weight=self.class_weights)
        else:
            loss_fct = nn.CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

class MetricsLoggerCallback(TrainerCallback):
    def __init__(self, log_file_path):
        self.log_file_path = log_file_path
        self.metrics_history = []

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs:
            entry = {"step": state.global_step, "epoch": state.epoch}
            entry.update(logs)
            self.metrics_history.append(entry)
            with open(self.log_file_path, "w") as f:
                json.dump(self.metrics_history, f, indent=2)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    acc = accuracy_score(labels, preds)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0)
    return {
        "accuracy": acc,
        "precision": precision,
        "recall": recall,
        "f1": f1
    }

def evaluate_threshold_sweep(logits, true_labels):
    """
    Sweeps confidence thresholds tau in [0.50 .. 0.85].
    If max(P(positive_edge)) < tau, defaults prediction to NONE (3).
    Calculates Positive Edge Precision and False Positive Edge Rate.
    """
    probs = torch.softmax(torch.tensor(logits), dim=-1).numpy()
    
    print("\n" + "="*80)
    print("🎯 CONSERVATIVE THRESHOLD SWEEP FOR GRAPH PURITY (DEFAULT TO 'NONE')")
    print("="*80)
    print(f"{'Tau':<8} | {'Overall Acc':<12} | {'Pos Edge Precision':<20} | {'FP Edge Rate':<18} | {'NONE Recall':<12}")
    print("-" * 80)

    best_tau = 0.50
    best_precision = 0.0

    results = []
    for tau in np.arange(0.50, 0.86, 0.05):
        thresholded_preds = []
        for i in range(len(probs)):
            p = probs[i]
            pos_probs = p[:3]  # SHAPES, DEPENDS_ON, CONFLICTS_WITH
            max_pos_idx = np.argmax(pos_probs)
            max_pos_prob = pos_probs[max_pos_idx]

            if max_pos_prob >= tau:
                thresholded_preds.append(max_pos_idx)
            else:
                thresholded_preds.append(3)  # Default to NONE

        thresholded_preds = np.array(thresholded_preds)
        acc = accuracy_score(true_labels, thresholded_preds)

        # Positive Edge Precision: Of all predicted positive edges (0, 1, 2), how many were actually correct?
        pos_mask = (thresholded_preds < 3)
        if np.sum(pos_mask) > 0:
            pos_precision = accuracy_score(true_labels[pos_mask], thresholded_preds[pos_mask])
        else:
            pos_precision = 1.0

        # False Positive Edge Rate: Of all actual NONE pairs (3), how many were incorrectly assigned a positive edge?
        actual_none_mask = (true_labels == 3)
        if np.sum(actual_none_mask) > 0:
            fp_edge_rate = np.mean(thresholded_preds[actual_none_mask] < 3)
            none_recall = np.mean(thresholded_preds[actual_none_mask] == 3)
        else:
            fp_edge_rate = 0.0
            none_recall = 1.0

        print(f"τ = {tau:.2f}  | {acc*100:6.2f}%      | {pos_precision*100:14.2f}%       | {fp_edge_rate*100:12.2f}%       | {none_recall*100:8.2f}%")
        
        results.append({
            "tau": float(tau),
            "accuracy": float(acc),
            "positive_edge_precision": float(pos_precision),
            "false_positive_edge_rate": float(fp_edge_rate),
            "none_recall": float(none_recall)
        })

        if pos_precision >= 0.88 and pos_precision > best_precision:
            best_precision = pos_precision
            best_tau = float(tau)

    print("="*80)
    print(f"💡 Recommended Graph Conservation Threshold: τ = {best_tau:.2f} (Pos Edge Precision: {best_precision*100:.2f}%)\n")
    return results, best_tau

class DatasetDictWrapper(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item["labels"] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

def main():
    print("🚀 Starting Vox v7 Gate 3 6-Epoch ModernBERT Fine-Tuning Pipeline (Pure CPU Mode)...")
    
    # 1. Load Dataset
    if not os.path.exists(DATASET_PATH):
        raise FileNotFoundError(f"Dataset file not found at {DATASET_PATH}")
        
    with open(DATASET_PATH, 'r') as f:
        data = json.load(f)
    pairs = data.get('pairs', data)
    print(f"Loaded {len(pairs)} ground-truth pairs from {DATASET_PATH}")

    # 2. Extract Texts, Labels, and Group Keys (Fact A)
    texts = [format_input(p['fact_a'], p['fact_b'], p['context']) for p in pairs]
    labels = [LABEL2ID[p['expected_label']] for p in pairs]
    groups = [p['fact_a'].strip().lower() for p in pairs]

    # 3. Stratified Group Split (80% Train / 10% Val / 10% Test) — Zero Fact Leakage
    np.random.seed(42)
    indices = np.arange(len(pairs))
    
    sgkf1 = StratifiedGroupKFold(n_splits=10)
    train_val_idx, test_idx = next(sgkf1.split(indices, labels, groups))
    
    train_val_labels = [labels[i] for i in train_val_idx]
    train_val_groups = [groups[i] for i in train_val_idx]
    sgkf2 = StratifiedGroupKFold(n_splits=9)
    train_sub_idx, val_sub_idx = next(sgkf2.split(train_val_idx, train_val_labels, train_val_groups))
    
    train_idx = train_val_idx[train_sub_idx]
    val_idx = train_val_idx[val_sub_idx]

    # Verify zero group leakage between Train and Test
    train_groups = set(groups[i] for i in train_idx)
    test_groups = set(groups[i] for i in test_idx)
    overlap = train_groups.intersection(test_groups)
    print(f"Dataset Splits -> Train: {len(train_idx)} | Val: {len(val_idx)} | Test: {len(test_idx)}")
    print(f"Fact Group Leakage Count between Train & Test: {len(overlap)} (MUST BE 0)")

    # Inverse Class Frequency Weights
    train_labels = [labels[i] for i in train_idx]
    label_counts = Counter(train_labels)
    total_samples = len(train_labels)
    num_classes = 4
    class_weights = [total_samples / (num_classes * label_counts[i]) for i in range(num_classes)]

    # Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    def encode_split(split_idx):
        split_texts = [texts[i] for i in split_idx]
        split_labels = [labels[i] for i in split_idx]
        encodings = tokenizer(split_texts, truncation=True, max_length=256, padding=False)
        return DatasetDictWrapper(encodings, split_labels)

    train_dataset = encode_split(train_idx)
    val_dataset = encode_split(val_idx)
    test_dataset = encode_split(test_idx)

    # ---------------------------------------------------------------------
    # PHASE 1: Fine-Tuning ModernBERT-base for 6 Epochs
    # ---------------------------------------------------------------------
    print("🏋️ PHASE 1: Fine-Tuning ModernBERT-base for 6 Epochs on Pure CPU...")
    
    fine_tune_model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME,
        num_labels=4,
        id2label=ID2LABEL,
        label2id=LABEL2ID
    )

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        eval_strategy="steps",
        eval_steps=50,
        save_strategy="steps",
        save_steps=50,
        save_total_limit=2,
        learning_rate=3e-5,
        per_device_train_batch_size=32,
        gradient_accumulation_steps=1,
        per_device_eval_batch_size=64,
        num_train_epochs=6,
        lr_scheduler_type="cosine",
        weight_decay=0.01,
        warmup_ratio=0.10,
        logging_steps=25,
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        use_cpu=True,
        report_to="none"
    )

    metrics_callback = MetricsLoggerCallback(METRICS_LOG_PATH)

    trainer = WeightedLossTrainer(
        class_weights=class_weights,
        model=fine_tune_model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        tokenizer=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=compute_metrics,
        callbacks=[metrics_callback]
    )

    trainer.train()

    best_model_dir = os.path.join(OUTPUT_DIR, "best_model")
    trainer.save_model(best_model_dir)
    tokenizer.save_pretrained(best_model_dir)
    print(f"Saved best 6-epoch PyTorch model to {best_model_dir}")

    # ---------------------------------------------------------------------
    # PHASE 2: Fine-Tuned Model Test Evaluation & Threshold Sweep
    # ---------------------------------------------------------------------
    print("\n🎯 PHASE 2: Evaluating 6-Epoch Model & Sweeping Confidence Thresholds...")
    ft_eval = trainer.predict(test_dataset)
    test_true_labels = np.array([labels[i] for i in test_idx])
    
    sweep_results, recommended_tau = evaluate_threshold_sweep(ft_eval.predictions, test_true_labels)

    # Save Calibration Results to JSON
    calib_json_path = os.path.expanduser("~/.vox/sandbox/threshold_calibration.json")
    with open(calib_json_path, "w") as f:
        json.dump({"sweep": sweep_results, "recommended_tau": recommended_tau}, f, indent=2)

    # ---------------------------------------------------------------------
    # PHASE 3: Export to INT8 Dynamic ONNX (model_quantized_v2.onnx)
    # ---------------------------------------------------------------------
    print("📦 PHASE 3: Exporting Fine-Tuned Model to INT8 Dynamic ONNX...")
    
    from optimum.onnxruntime import ORTModelForSequenceClassification, ORTQuantizer
    from optimum.onnxruntime.configuration import AutoQuantizationConfig

    onnx_temp_dir = os.path.join(OUTPUT_DIR, "onnx_temp")
    
    print("  Exporting PyTorch weights to FP32 ONNX...")
    ort_model = ORTModelForSequenceClassification.from_pretrained(best_model_dir, export=True)
    ort_model.save_pretrained(onnx_temp_dir)

    print("  Applying Dynamic INT8 Quantization (qint8)...")
    quantizer = ORTQuantizer.from_pretrained(onnx_temp_dir)
    dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)
    
    quantized_dir = os.path.join(OUTPUT_DIR, "onnx_quantized")
    quantizer.quantize(save_dir=quantized_dir, quantization_config=dqconfig)

    source_onnx = os.path.join(quantized_dir, "model_quantized.onnx")
    if not os.path.exists(source_onnx):
        source_onnx = os.path.join(quantized_dir, "model.onnx")
        
    shutil.copy(source_onnx, ONNX_EXPORT_PATH)
    file_size_mb = os.path.getsize(ONNX_EXPORT_PATH) / (1024 * 1024)
    print(f"✅ Exported 6-Epoch INT8 ONNX Model to {ONNX_EXPORT_PATH} ({file_size_mb:.2f} MB)")

    # ---------------------------------------------------------------------
    # PHASE 4: Proactive Disk Cleanup
    # ---------------------------------------------------------------------
    print("\n🧹 PHASE 4: Executing Disk Cleanup...")
    if os.path.exists(OUTPUT_DIR):
        shutil.rmtree(OUTPUT_DIR)
        print(f"  Removed temporary training dir: {OUTPUT_DIR}")
        
    hf_cache_dir = os.path.expanduser("~/.cache/huggingface")
    if os.path.exists(hf_cache_dir):
        shutil.rmtree(hf_cache_dir)
        print(f"  Cleared HuggingFace download cache: {hf_cache_dir}")

    df_output = subprocess.check_output(["df", "-h", os.path.expanduser("~")]).decode("utf-8")
    print(f"\nServer Storage Status:\n{df_output}")
    print("🎉 6-Epoch Pipeline execution complete!")

if __name__ == "__main__":
    main()