"""
Tahap 9: Testing Abjad SIBI

Sesuai README.
"""

import sys
import json
import numpy as np
from pathlib import Path
import argparse

import torch
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, precision_recall_fscore_support
import matplotlib.pyplot as plt
import seaborn as sns

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

from utils.config import TEST_DIR, MODEL_DIR, BATCH_SIZE, OUTPUT_DIR, CONFUSION_MATRIX_DIR, CLASSIFICATION_REPORT_DIR


class AbjadDataset:
    def __init__(self, data_dir: Path):
        self.data_dir = data_dir
        self.samples = []
        self.labels = []
        self.label_to_idx = {}
        self._load_data()

    def _load_data(self):
        if not self.data_dir.exists():
            return
        label_dirs = sorted([d for d in self.data_dir.iterdir() if d.is_dir()])
        self.label_to_idx = {d.name: idx for idx, d in enumerate(label_dirs)}
        for label_dir in label_dirs:
            label_idx = self.label_to_idx[label_dir.name]
            for npy_file in label_dir.glob("*.npy"):
                try:
                    self.samples.append(np.load(npy_file))
                    self.labels.append(label_idx)
                except:
                    pass

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


def plot_confusion_matrix(cm, labels, output_path, title):
    plt.figure(figsize=(14, 12))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=labels, yticklabels=labels)
    plt.title(title)
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    plt.savefig(output_path, dpi=150)
    plt.close()
    print(f"Confusion matrix saved to: {output_path}")


def test_abjad_sibi(data_dir: Path, model_dir: Path, output_dir: Path):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")

    model_path = model_dir / "abjad_sibi" / "model_abjad_sibi.pt"
    encoder_path = model_dir / "abjad_sibi" / "label_encoder.json"

    if not model_path.exists():
        print(f"Error: Model tidak ditemukan di {model_path}")
        return

    print(f"Loading model from: {model_path}")
    model = torch.load(model_path, map_location=device, weights_only=False)
    model.eval()

    with open(encoder_path, 'r') as f:
        idx_to_label = json.load(f)["idx_to_label"]

    print(f"Loading test dataset...")
    dataset = AbjadDataset(data_dir / "test" / "abjad_sibi")
    print(f"Test samples: {len(dataset)}")

    if len(dataset) == 0:
        print("No test samples found!")
        return

    sequences = torch.FloatTensor(np.array(dataset.samples))
    true_labels = np.array(dataset.labels)

    print("Generating predictions...")
    predictions = []
    with torch.no_grad():
        for i in range(0, len(sequences), BATCH_SIZE):
            batch = sequences[i:i+BATCH_SIZE].to(device)
            outputs = model(batch)
            _, preds = torch.max(outputs, 1)
            predictions.extend(preds.cpu().numpy().tolist())

    accuracy = accuracy_score(true_labels, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(true_labels, predictions, average='weighted', zero_division=0)
    cm = confusion_matrix(true_labels, predictions)

    label_names = [idx_to_label[str(i)] for i in range(len(idx_to_label))]
    report = classification_report(true_labels, predictions, target_names=label_names, zero_division=0)

    print("\n" + "="*60)
    print("TESTING ABJAD SIBI MODEL")
    print("="*60)
    print(f"\nAccuracy:  {accuracy*100:.2f}%")
    print(f"Precision: {precision*100:.2f}%")
    print(f"Recall:    {recall*100:.2f}%")
    print(f"F1-Score: {f1*100:.2f}%")
    print("\n" + report)

    plot_confusion_matrix(cm, label_names, CONFUSION_MATRIX_DIR / "abjad_sibi_confusion_matrix.png", "Abjad SIBI Confusion Matrix")

    with open(CLASSIFICATION_REPORT_DIR / "abjad_sibi_classification_report.txt", 'w') as f:
        f.write("Classification Report - Abjad SIBI\n")
        f.write("="*60 + "\n\n")
        f.write(f"Accuracy: {accuracy*100:.2f}%\n")
        f.write(f"Precision: {precision*100:.2f}%\n")
        f.write(f"Recall: {recall*100:.2f}%\n")
        f.write(f"F1-Score: {f1*100:.2f}%\n\n")
        f.write(report)

    results = {"accuracy": float(accuracy), "precision": float(precision), "recall": float(recall), "f1_score": float(f1)}
    with open(output_dir / "abjad_sibi_test_results.json", 'w') as f:
        json.dump(results, f, indent=2)

    print("\n" + "="*60)
    print("TESTING COMPLETE")
    print("="*60)


def main():
    parser = argparse.ArgumentParser(description="Testing model abjad SIBI")
    parser.add_argument("--data-dir", type=str, default=None)
    parser.add_argument("--model-dir", type=str, default=None)
    parser.add_argument("--output-dir", type=str, default=None)
    args = parser.parse_args()

    data_dir = Path(args.data_dir) if args.data_dir else TEST_DIR.parent
    model_dir = Path(args.model_dir) if args.model_dir else MODEL_DIR
    output_dir = Path(args.output_dir) if args.output_dir else OUTPUT_DIR

    print("="*60)
    print("TAHAP 9: TESTING ABJAD SIBI")
    print("="*60)

    test_abjad_sibi(data_dir, model_dir, output_dir)


if __name__ == "__main__":
    main()