import os
import shutil
import tempfile

import pickle
import pandas as pd
from typing import Tuple


from image_embedding_model.imagebind_embedding import calc_imagebind_embedding
from image_embedding_model.openclip_embedding import calc_openclip_embedding
from brain_embedding_model.brain_embedding import calc_brain_embedding

BRAIN_EMBEDDING_MAT_IMAGEBIND = (
    "./brain_embedding_model/imagebind_embedding_model.safetensors"
)


TRANSFORMATION_TABLE_PATH = "./brain_embedding_model/va_transformation_table.csv"

## The below two model are generated by executing the notebook open_brain_examples.ipynb
VALENCE_MODEL_PATH = "./valence_prediction.pickle"
AROUSAL_MODEL_PATH = "./arousal_prediction.pickle"

MAX_PREDICTED_VALUE = 7
MIN_PREDICTED_VALUE = 1

def copy_image_to_temp(src_image_path: str) -> str:
    """
    Copies the image at src_image_path into a newly created temporary directory.

    Args:
        src_image_path (str): The file path of the source image.

    Returns:
        str: The file path of the copied image in the temporary directory.
    """
    # Create a new temporary directory
    temp_dir = tempfile.mkdtemp(prefix="img_copy_")

    # Ensure the source file exists
    if not os.path.isfile(src_image_path):
        raise FileNotFoundError(f"Source image not found: {src_image_path}")

    # Define the destination path in the temporary directory
    image_name = os.path.basename(src_image_path)
    dest_image_path = os.path.join(temp_dir, image_name)

    # Copy the image
    shutil.copy2(src_image_path, dest_image_path)

    return temp_dir

def scale_values(value: float, min_range: float, max_range: float) -> float:
    """
    Scales an input value from the range [min_range, max_range] to the range [-1, 1]
    to match the valence-arousal space.

    Args:
        value (float): The value to scale.
        min_val (float): The minimum value of the original range.
        max_val (float): The maximum value of the original range.
    Returns:
        float: The scaled value.
    """
    # the target range is [-1, 1]
    TARGET_MIN = -1
    TARGET_MAX = 1
    return TARGET_MIN + (TARGET_MAX - TARGET_MIN)* ((value - min_range) / (max_range-min_range))


def get_closest_emotion(
    valence: float, arousal: float, emotions_df: pd.DataFrame
) -> Tuple[str, str, str]:
    """
    Find the closest emotion in the 2D valence–arousal space.

    Args:
        valence (float): Valence value to classify.
        arousal (float): Arousal value to classify.

    Returns:
        Tuple[str, str, str]: (emotion name, color name, hex code)
    """
    best = None
    best_dist = float("inf")

    # Scale valence and arousal to the range [-1, 1]
    scaled_valence = scale_values(valence, min_range= MIN_PREDICTED_VALUE, max_range= MAX_PREDICTED_VALUE)
    scaled_arousal = scale_values(arousal, min_range= MIN_PREDICTED_VALUE, max_range= MAX_PREDICTED_VALUE)

    for index, row in emotions_df.iterrows():
        # Euclidean distance squared
        dv = scaled_valence - row["Valence"]
        da = scaled_arousal - row["Arousal"]
        dist = dv * dv + da * da
        if dist < best_dist:
            best_dist = dist
            best = row
    return best["Emotion"]


def predict_emotion(in_image_path: str) -> Tuple[str, float, float]:
    """
    Predicts the closest emotion label, valence, and arousal scores for a given image.

    This function:
        1. Loads a valence–arousal transformation table.
        2. Copies the input image to a temporary folder.
        3. Calculates image embeddings and brain embeddings.
        4. Loads pretrained valence and arousal models.
        5. Predicts valence and arousal scores.
        6. Determines the closest emotion based on the scores.
        7. Cleans up the temporary folder.

    Parameters
    ----------
    in_image_path : str
        Path to the input image file.

    Returns
    -------
    Tuple[str, float, float]
        - emotion: Closest emotion label.
        - valence_pred: Predicted valence score.
        - arousal_pred: Predicted arousal score.
    """
    # Load the transformation table
    try:
        conversion_table = pd.read_csv(TRANSFORMATION_TABLE_PATH)
    except Exception as e:
        raise RuntimeError("Failed to load transformation table.")

    temp_folder = None

    try:
        # Copy image to temporary folder
        temp_folder = copy_image_to_temp(in_image_path)

        # Calculate embeddings
        image_emb = calc_imagebind_embedding(temp_folder)
        brain_emb = calc_brain_embedding(
            image_emb,
            brain_model_file=BRAIN_EMBEDDING_MAT_IMAGEBIND,
        )

        # Load models
        try:
            with open(VALENCE_MODEL_PATH, "rb") as f:
                valence_model = pickle.load(f)
        except Exception as e:
            raise RuntimeError("Failed to load valence model.")

        try:
            with open(AROUSAL_MODEL_PATH, "rb") as f:
                arousal_model = pickle.load(f)
        except Exception as e:
            raise RuntimeError("Failed to load arousal model.")

        # Prepare data for prediction
        try:
            embedding = brain_emb[0]["brain_embedding"].reshape(1, -1)
            valence_pred = float(valence_model.predict(embedding)[0])
            arousal_pred = float(arousal_model.predict(embedding)[0])
        except Exception as e:
            raise RuntimeError("Error during valence/arousal prediction.")

        # Determine emotion
        emotion = get_closest_emotion(valence_pred, arousal_pred, conversion_table)
        return emotion, valence_pred, arousal_pred

    finally:
        # Cleanup temporary folder
        if temp_folder and os.path.isdir(temp_folder):
            shutil.rmtree(temp_folder)


def call_example():
    in_image_path = "./examples/sample_OASIS_images/Beach 1.jpg"
    emotion, valence_pred, arousal_pred = predict_emotion(in_image_path)
    print(f"Emotion: {emotion}, Valence:{valence_pred}, Arousal:{arousal_pred}")


# call_example()
