"""Evaluation script for Myanmar Ghost sentiment model."""

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any, Dict

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from src.utils.logger import setup_logger
from src.utils.metrics import compute_metrics, compute_confusion_matrix

logger = setup_logger("evaluate", log_dir="outputs/logs")


class SentimentDataset(Dataset):
    """Dataset for sentiment classification."""
    
    def __init__(
        self,
        data,
        tokenizer,
        max_length: int = 512,
        label_mapping: dict = None,
    ):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.label_mapping = label_mapping or {
            "negative": 0, "neutral": 1, "positive": 2, "sarcastic": 3
        }
    
    def __len__(self) -> int:
        return len(self.data)
    
    def __getitem__(self, idx: int):
        item = self.data[idx]
        
        encoding = self.tokenizer(
            item["text"],
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )
        
        label = self.label_mapping.get(item.get("label", "neutral"), 1)
        
        return (
            encoding["input_ids"].squeeze(0),
            encoding["attention_mask"].squeeze(0),
            torch.tensor(label, dtype=torch.long),
            item.get("text", ""),
        )


def load_data(data_path: str):
    """Load evaluation data."""
    if data_path.endswith(".jsonl"):
        data = []
        with open(data_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    data.append(json.loads(line))
    elif data_path.endswith(".json"):
        with open(data_path, "r", encoding="utf-8") as f:
            data = json.load(f)
    else:
        raise ValueError(f"Unsupported format: {data_path}")
    
    return data


def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    device: torch.device,
    class_names: list = None,
) -> Dict[str, Any]:
    """Evaluate the model."""
    if class_names is None:
        class_names = ["negative", "neutral", "positive", "sarcastic"]
    
    model.eval()
    
    all_predictions = []
    all_labels = []
    all_texts = []
    all_probabilities = []
    
    with torch.no_grad():
        for input_ids, attention_mask, labels, texts in tqdm(dataloader, desc="Evaluating"):
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            
            outputs = model(input_ids, attention_mask)
            probs = torch.softmax(outputs, dim=-1)
            
            predictions = outputs.argmax(dim=-1).cpu().tolist()
            all_predictions.extend(predictions)
            all_labels.extend(labels.tolist())
            all_texts.extend(texts)
            all_probabilities.extend(probs.cpu().numpy().tolist())
    
    # Compute metrics
    metrics = compute_metrics(all_predictions, all_labels, class_names)
    
    # Confusion matrix
    cm = compute_confusion_matrix(all_predictions, all_labels)
    
    # Per-sample results
    results = []
    for i, (text, label, pred, probs) in enumerate(zip(
        all_texts, all_labels, all_predictions, all_probabilities
    )):
        results.append({
            "text": text,
            "true_label": class_names[label],
            "predicted_label": class_names[pred],
            "correct": label == pred,
            "probabilities": {
                class_names[j]: probs[j] for j in range(len(class_names))
            },
        })
    
    return {
        "metrics": metrics,
        "confusion_matrix": cm.tolist(),
        "class_names": class_names,
        "results": results,
    }


def save_results(results: Dict, output_path: str) -> None:
    """Save evaluation results to file."""
    output_dir = Path(output_path).parent
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Save full results
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    # Save summary
    summary_path = output_dir / "evaluation_summary.txt"
    with open(summary_path, "w", encoding="utf-8") as f:
        f.write("=" * 60 + "\n")
        f.write("EVALUATION SUMMARY\n")
        f.write("=" * 60 + "\n\n")
        
        metrics = results["metrics"]
        f.write(f"Accuracy:      {metrics['accuracy']:.4f}\n")
        f.write(f"F1 (weighted): {metrics['f1_weighted']:.4f}\n")
        f.write(f"F1 (macro):    {metrics['f1_macro']:.4f}\n")
        f.write(f"Precision:     {metrics['precision_weighted']:.4f}\n")
        f.write(f"Recall:        {metrics['recall_weighted']:.4f}\n\n")
        
        f.write("Per-class F1:\n")
        for name in results["class_names"]:
            f1_key = f"f1_{name}"
            if f1_key in metrics:
                f.write(f"  {name}: {metrics[f1_key]:.4f}\n")
        
        f.write(f"\nConfusion Matrix:\n")
        f.write(str(np.array(results["confusion_matrix"])) + "\n")
    
    logger.info(f"Results saved to {output_path}")
    logger.info(f"Summary saved to {summary_path}")


def main(args):
    """Main evaluation function."""
    logger.info("Starting evaluation...")
    logger.info(f"Arguments: {vars(args)}")
    
    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
    logger.info(f"Using device: {device}")
    
    # Load tokenizer
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    
    # Load data
    logger.info(f"Loading data from {args.data_path}")
    data = load_data(args.data_path)
    logger.info(f"Total samples: {len(data)}")
    
    # Create dataset and dataloader
    dataset = SentimentDataset(data, tokenizer, args.max_length)
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=2,
    )
    
    # Load model
    logger.info(f"Loading model from {args.model_path}")
    from src.models.transformer_model import TransformerSentimentModel
    
    model = TransformerSentimentModel(
        model_name=args.model_name,
        num_labels=4,
    )
    
    checkpoint = torch.load(args.model_path, map_location=device)
    if "model_state_dict" in checkpoint:
        model.load_state_dict(checkpoint["model_state_dict"])
    else:
        model.load_state_dict(checkpoint)
    
    model.to(device)
    
    # Evaluate
    results = evaluate(model, dataloader, device)
    
    # Print summary
    logger.info("\n" + "=" * 60)
    logger.info("EVALUATION RESULTS")
    logger.info("=" * 60)
    logger.info(f"Accuracy:      {results['metrics']['accuracy']:.4f}")
    logger.info(f"F1 (weighted): {results['metrics']['f1_weighted']:.4f}")
    logger.info(f"Precision:     {results['metrics']['precision_weighted']:.4f}")
    logger.info(f"Recall:        {results['metrics']['recall_weighted']:.4f}")
    
    # Save results
    output_path = args.output_path or "outputs/results/evaluation_results.json"
    save_results(results, output_path)
    
    return results["metrics"]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate Myanmar Ghost model")
    
    parser.add_argument("--data_path", type=str, required=True, help="Test data file")
    parser.add_argument("--model_path", type=str, required=True, help="Model checkpoint")
    parser.add_argument("--model_name", type=str, default="bert-base-multilingual-cased")
    parser.add_argument("--output_path", type=str, default=None, help="Output path")
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--max_length", type=int, default=512)
    parser.add_argument("--cpu", action="store_true", help="Use CPU only")
    
    args = parser.parse_args()
    main(args)