import inspect
import json
import os
import random
import re
import unicodedata
from typing import Dict

import numpy as np
import pandas as pd
import torch

from datasets import DatasetDict, load_dataset
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    precision_recall_fscore_support,
)
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)


# =========================================================
# CONFIGURATION
# =========================================================

MODEL_NAME = os.getenv("MODEL_NAME", "xlm-roberta-base")
DATASET_REPO = os.getenv(
    "DATASET_REPO", "d12o6aa/ArabGuard-Egyptian-V1"
)

MODEL_OUTPUT_DIR = os.getenv("MODEL_OUTPUT_DIR", "./arabguard_model")
CHECKPOINT_DIR = os.getenv("CHECKPOINT_DIR", "./arabguard_checkpoints")
DASHBOARD_DATA_DIR = os.getenv("DASHBOARD_DATA_DIR", "./dashboard_data")

MAX_LENGTH = int(os.getenv("MAX_LENGTH", "128"))
NUM_EPOCHS = float(os.getenv("NUM_EPOCHS", "4"))
LEARNING_RATE = float(os.getenv("LEARNING_RATE", "2e-5"))
TRAIN_BATCH_SIZE = int(os.getenv("TRAIN_BATCH_SIZE", "8"))
EVAL_BATCH_SIZE = int(os.getenv("EVAL_BATCH_SIZE", "8"))
RANDOM_SEED = 42


# =========================================================
# REPRODUCIBILITY
# =========================================================

def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


set_seed(RANDOM_SEED)


# =========================================================
# TEXT NORMALIZATION
# =========================================================

def remove_arabic_diacritics(text: str) -> str:
    arabic_diacritics = re.compile(
        r"""
        ّ    |
        َ    |
        ً    |
        ُ    |
        ٌ    |
        ِ    |
        ٍ    |
        ْ    |
        ـ
        """,
        re.VERBOSE,
    )

    return re.sub(arabic_diacritics, "", text)


def normalize_arabic_letters(text: str) -> str:
    replacements = {
        "أ": "ا",
        "إ": "ا",
        "آ": "ا",
        "ٱ": "ا",
        "ى": "ي",
        "ؤ": "و",
        "ئ": "ي",
    }

    for old, new in replacements.items():
        text = text.replace(old, new)

    return text


def normalize_text(text: str) -> str:
    if text is None:
        return ""

    text = str(text)

    # Normalize Unicode representations.
    text = unicodedata.normalize("NFKC", text)

    # Remove zero-width and direction control characters.
    text = re.sub(
        r"[\u200B-\u200F\u202A-\u202E\u2060-\u206F\uFEFF]",
        "",
        text,
    )

    text = remove_arabic_diacritics(text)
    text = normalize_arabic_letters(text)

    # Normalize URLs, emails and long numbers.
    text = re.sub(
        r"https?://\S+|www\.\S+",
        " URL ",
        text,
        flags=re.IGNORECASE,
    )

    text = re.sub(
        r"\b[\w.\-+]+@[\w.\-]+\.\w+\b",
        " EMAIL ",
        text,
        flags=re.IGNORECASE,
    )

    text = re.sub(
        r"\b\d{5,}\b",
        " NUMBER ",
        text,
    )

    # Reduce exaggerated repeated characters.
    text = re.sub(
        r"(.)\1{4,}",
        r"\1\1",
        text,
    )

    # Reduce repeated punctuation.
    text = re.sub(
        r"([!?.,،؛:])\1+",
        r"\1",
        text,
    )

    # Normalize whitespace.
    text = re.sub(
        r"\s+",
        " ",
        text,
    ).strip()

    return text


# =========================================================
# LOAD DATASET FILES SEPARATELY
# =========================================================

print("Loading dataset files...")

train_dataset = load_dataset(
    "csv",
    data_files=f"hf://datasets/{DATASET_REPO}/train.csv",
    split="train",
)

validation_dataset = load_dataset(
    "csv",
    data_files=f"hf://datasets/{DATASET_REPO}/val.csv",
    split="train",
)

test_dataset = load_dataset(
    "csv",
    data_files=f"hf://datasets/{DATASET_REPO}/test.csv",
    split="train",
)

dataset = DatasetDict(
    {
        "train": train_dataset,
        "validation": validation_dataset,
        "test": test_dataset,
    }
)

print(dataset)


# =========================================================
# CLEAN DATASET
# =========================================================

def clean_example(example: Dict) -> Dict:
    text = str(example.get("text", "")).strip()
    label = int(example.get("label", 0))

    return {
        "text": text,
        "normalized_text": normalize_text(text),
        "label": label,
    }


dataset = dataset.map(clean_example)

for split_name in dataset.keys():
    columns_to_remove = [
        column
        for column in dataset[split_name].column_names
        if column not in ["text", "normalized_text", "label"]
    ]

    if columns_to_remove:
        dataset[split_name] = dataset[split_name].remove_columns(
            columns_to_remove
        )


# Remove empty rows.
def valid_example(example: Dict) -> bool:
    return bool(example["text"].strip())


dataset = dataset.filter(valid_example)

print("\nCleaned dataset:")
print(dataset)
print("\nExample:")
print(dataset["train"][0])


# =========================================================
# LABEL CONFIGURATION
# =========================================================

unique_labels = sorted(
    set(dataset["train"]["label"])
)

label_names = [
    str(label)
    for label in unique_labels
]

label2id = {
    label_name: index
    for index, label_name in enumerate(label_names)
}

id2label = {
    index: label_name
    for index, label_name in enumerate(label_names)
}

print("\nLabel mappings:")
print("label2id:", label2id)
print("id2label:", id2label)


def encode_label(example: Dict) -> Dict:
    example["labels"] = label2id[
        str(example["label"])
    ]

    return example


dataset = dataset.map(encode_label)


# =========================================================
# TOKENIZER
# =========================================================

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME
)


def tokenize_normalized_batch(batch: Dict) -> Dict:
    return tokenizer(
        batch["normalized_text"],
        truncation=True,
        max_length=MAX_LENGTH,
    )


tokenized_dataset = dataset.map(
    tokenize_normalized_batch,
    batched=True,
)

for split_name in tokenized_dataset.keys():
    columns_to_remove = [
        column
        for column in tokenized_dataset[split_name].column_names
        if column not in [
            "input_ids",
            "attention_mask",
            "labels",
        ]
    ]

    if columns_to_remove:
        tokenized_dataset[split_name] = (
            tokenized_dataset[split_name]
            .remove_columns(columns_to_remove)
        )


data_collator = DataCollatorWithPadding(
    tokenizer=tokenizer
)


# =========================================================
# MODEL
# =========================================================

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=len(label_names),
    id2label=id2label,
    label2id=label2id,
)


# =========================================================
# METRICS
# =========================================================

def calculate_metrics_from_arrays(
    labels: np.ndarray,
    predictions: np.ndarray,
) -> Dict[str, float]:
    precision, recall, f1, _ = (
        precision_recall_fscore_support(
            labels,
            predictions,
            average="weighted",
            zero_division=0,
        )
    )

    accuracy = accuracy_score(
        labels,
        predictions,
    )

    return {
        "accuracy": float(accuracy),
        "precision": float(precision),
        "recall": float(recall),
        "f1": float(f1),
    }


def compute_metrics(eval_prediction) -> Dict[str, float]:
    logits, labels = eval_prediction

    predictions = np.argmax(
        logits,
        axis=-1,
    )

    return calculate_metrics_from_arrays(
        labels,
        predictions,
    )


# =========================================================
# TRAINING ARGUMENTS
# =========================================================

training_argument_parameters = inspect.signature(
    TrainingArguments.__init__
).parameters

training_arguments_dictionary = {
    "output_dir": CHECKPOINT_DIR,
    "learning_rate": LEARNING_RATE,
    "num_train_epochs": NUM_EPOCHS,
    "per_device_train_batch_size": TRAIN_BATCH_SIZE,
    "per_device_eval_batch_size": EVAL_BATCH_SIZE,
    "weight_decay": 0.01,
    "save_strategy": "epoch",
    "logging_strategy": "steps",
    "logging_steps": 20,
    "load_best_model_at_end": True,
    "metric_for_best_model": "f1",
    "greater_is_better": True,
    "save_total_limit": 2,
    "report_to": "none",
    "fp16": torch.cuda.is_available(),
    "seed": RANDOM_SEED,
    "data_seed": RANDOM_SEED,
}

if "eval_strategy" in training_argument_parameters:
    training_arguments_dictionary[
        "eval_strategy"
    ] = "epoch"

elif "evaluation_strategy" in training_argument_parameters:
    training_arguments_dictionary[
        "evaluation_strategy"
    ] = "epoch"

training_arguments = TrainingArguments(
    **training_arguments_dictionary
)


# =========================================================
# TRAINER
# =========================================================

trainer_arguments = {
    "model": model,
    "args": training_arguments,
    "train_dataset": tokenized_dataset["train"],
    "eval_dataset": tokenized_dataset["validation"],
    "data_collator": data_collator,
    "compute_metrics": compute_metrics,
}

trainer_signature = inspect.signature(
    Trainer.__init__
).parameters

if "processing_class" in trainer_signature:
    trainer_arguments["processing_class"] = tokenizer

elif "tokenizer" in trainer_signature:
    trainer_arguments["tokenizer"] = tokenizer

trainer = Trainer(
    **trainer_arguments
)


# =========================================================
# TRAIN MODEL
# =========================================================

print("\nTraining started...")
training_result = trainer.train()

print("\nTraining finished.")


# =========================================================
# EVALUATE NORMALIZED DATA
# =========================================================

validation_results = trainer.evaluate(
    tokenized_dataset["validation"],
    metric_key_prefix="validation",
)

normalized_test_output = trainer.predict(
    tokenized_dataset["test"]
)

normalized_predictions = np.argmax(
    normalized_test_output.predictions,
    axis=-1,
)

normalized_labels = normalized_test_output.label_ids

normalized_metrics = calculate_metrics_from_arrays(
    normalized_labels,
    normalized_predictions,
)


# =========================================================
# EVALUATE RAW DATA
# =========================================================

def create_raw_tokenized_test_dataset():
    raw_test_dataset = dataset["test"].map(
        lambda batch: tokenizer(
            batch["text"],
            truncation=True,
            max_length=MAX_LENGTH,
        ),
        batched=True,
    )

    columns_to_remove = [
        column
        for column in raw_test_dataset.column_names
        if column not in [
            "input_ids",
            "attention_mask",
            "labels",
        ]
    ]

    if columns_to_remove:
        raw_test_dataset = raw_test_dataset.remove_columns(
            columns_to_remove
        )

    return raw_test_dataset


raw_test_dataset = create_raw_tokenized_test_dataset()

raw_test_output = trainer.predict(
    raw_test_dataset
)

raw_predictions = np.argmax(
    raw_test_output.predictions,
    axis=-1,
)

raw_labels = raw_test_output.label_ids

raw_metrics = calculate_metrics_from_arrays(
    raw_labels,
    raw_predictions,
)


# =========================================================
# CONFUSION MATRIX
# =========================================================

matrix = confusion_matrix(
    normalized_labels,
    normalized_predictions,
    labels=list(range(len(label_names))),
)

confusion_matrix_dataframe = pd.DataFrame(
    matrix,
    index=[
        f"Actual {id2label[index]}"
        for index in range(len(label_names))
    ],
    columns=[
        f"Predicted {id2label[index]}"
        for index in range(len(label_names))
    ],
)


# =========================================================
# CLASSIFICATION REPORT
# =========================================================

classification_report_data = classification_report(
    normalized_labels,
    normalized_predictions,
    target_names=[
        id2label[index]
        for index in range(len(label_names))
    ],
    output_dict=True,
    zero_division=0,
)


# =========================================================
# SAVE MODEL
# =========================================================

os.makedirs(
    MODEL_OUTPUT_DIR,
    exist_ok=True,
)

trainer.save_model(
    MODEL_OUTPUT_DIR
)

tokenizer.save_pretrained(
    MODEL_OUTPUT_DIR
)


# =========================================================
# SAVE DASHBOARD DATA
# =========================================================

os.makedirs(
    DASHBOARD_DATA_DIR,
    exist_ok=True,
)

training_history = trainer.state.log_history

history_dataframe = pd.DataFrame(
    training_history
)

history_dataframe.to_csv(
    os.path.join(
        DASHBOARD_DATA_DIR,
        "training_history.csv",
    ),
    index=False,
)

confusion_matrix_dataframe.to_csv(
    os.path.join(
        DASHBOARD_DATA_DIR,
        "confusion_matrix.csv",
    ),
)

with open(
    os.path.join(
        DASHBOARD_DATA_DIR,
        "classification_report.json",
    ),
    "w",
    encoding="utf-8",
) as file:
    json.dump(
        classification_report_data,
        file,
        ensure_ascii=False,
        indent=4,
    )


normalization_accuracy_change = (
    normalized_metrics["accuracy"]
    - raw_metrics["accuracy"]
)

metrics_data = {
    "model_name": MODEL_NAME,
    "model_output_directory": MODEL_OUTPUT_DIR,
    "max_length": MAX_LENGTH,
    "epochs": NUM_EPOCHS,
    "device_used_for_training": (
        "cuda"
        if torch.cuda.is_available()
        else "cpu"
    ),
    "dataset": {
        "train_samples": len(dataset["train"]),
        "validation_samples": len(
            dataset["validation"]
        ),
        "test_samples": len(dataset["test"]),
    },
    "raw_test_metrics": raw_metrics,
    "normalized_test_metrics": normalized_metrics,
    "normalization_accuracy_change": float(
        normalization_accuracy_change
    ),
    "validation_metrics": {
        key: float(value)
        for key, value in validation_results.items()
        if isinstance(
            value,
            (
                int,
                float,
                np.integer,
                np.floating,
            ),
        )
    },
    "training_metrics": {
        key: float(value)
        for key, value in training_result.metrics.items()
        if isinstance(
            value,
            (
                int,
                float,
                np.integer,
                np.floating,
            ),
        )
    },
    "label_mapping": {
        str(key): value
        for key, value in id2label.items()
    },
}

with open(
    os.path.join(
        DASHBOARD_DATA_DIR,
        "metrics.json",
    ),
    "w",
    encoding="utf-8",
) as file:
    json.dump(
        metrics_data,
        file,
        ensure_ascii=False,
        indent=4,
    )


# =========================================================
# PRINT FINAL RESULTS
# =========================================================

print("\n" + "=" * 60)
print("RAW TEST METRICS")
print("=" * 60)

for metric_name, metric_value in raw_metrics.items():
    print(
        f"{metric_name}: "
        f"{metric_value:.4f}"
    )

print("\n" + "=" * 60)
print("NORMALIZED TEST METRICS")
print("=" * 60)

for metric_name, metric_value in normalized_metrics.items():
    print(
        f"{metric_name}: "
        f"{metric_value:.4f}"
    )

print("\nNormalization accuracy change:")

print(
    f"{normalization_accuracy_change:+.4f}"
)

print(
    f"\nModel saved to: "
    f"{MODEL_OUTPUT_DIR}"
)

print(
    f"Dashboard data saved to: "
    f"{DASHBOARD_DATA_DIR}"
)