import torch
import pandas as pd

from sentence_transformers import SentenceTransformer, util
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

PMFC_JSON_PATH = "assets/pmfc.json"

EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
CLASSIFIER_MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"


# =========================
# LAZY-LOADED MODEL / DATA CACHE
# =========================
_CACHE = {
    "embedder": None,
    "tokenizer": None,
    "model": None,
    "pipe": None,
    "pm_examples": None,    # DataFrame
    "fc_examples": None,    # DataFrame
    "pm_embeddings": None,
    "fc_embeddings": None,
}


def _get_embedder():
    if _CACHE["embedder"] is None:
        _CACHE["embedder"] = SentenceTransformer(EMBED_MODEL_NAME, device=str(device))
    return _CACHE["embedder"]


def _get_pipe():
    if _CACHE["pipe"] is None:
        tokenizer = AutoTokenizer.from_pretrained(CLASSIFIER_MODEL_NAME)
        model = AutoModelForCausalLM.from_pretrained(
            CLASSIFIER_MODEL_NAME,
            torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
            device_map="auto" if device.type == "cuda" else None
        )
        if device.type == "cpu":
            model.to(device)

        _CACHE["tokenizer"] = tokenizer
        _CACHE["model"] = model
        _CACHE["pipe"] = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            return_full_text=False,
            do_sample=False
        )
    return _CACHE["pipe"], _CACHE["tokenizer"]


def _load_pmfc_corpus():
    """Loads assets/pmfc.json and pre-computes PM/FC embeddings separately,
    same split logic as the reference script (str.contains on FinalAssessment)."""
    if _CACHE["pm_examples"] is not None:
        return

    examples_df = pd.read_json(PMFC_JSON_PATH)

    pm_examples = examples_df[
        examples_df["FinalAssessment"].str.contains("PM", na=False)
    ].reset_index(drop=True)

    fc_examples = examples_df[
        examples_df["FinalAssessment"].str.contains("FC", na=False)
    ].reset_index(drop=True)

    embedder = _get_embedder()

    pm_embeddings = embedder.encode(
        pm_examples["Paragraph_content"].tolist(),
        convert_to_tensor=True,
        normalize_embeddings=True
    )
    fc_embeddings = embedder.encode(
        fc_examples["Paragraph_content"].tolist(),
        convert_to_tensor=True,
        normalize_embeddings=True
    )

    _CACHE["pm_examples"] = pm_examples
    _CACHE["fc_examples"] = fc_examples
    _CACHE["pm_embeddings"] = pm_embeddings
    _CACHE["fc_embeddings"] = fc_embeddings


# =========================
# PROMPT (matches reference script's generate_pmfc_prompt exactly)
# =========================
def generate_pmfc_prompt(paragraph, pm_example, fc_example):
    prompt = f"""
You are an expert sustainability-report analyst.

The paragraph below is already known to be:
- Sustainability Relevant
- Quantitative

Your task is to classify it into EXACTLY ONE category.

Definitions

Performance Metrics (PM):
- Reports achievements
- Reports measured performance
- Reports historical results
- Reports current status
- Reports quantities already achieved

Future Commitments (FC):
- Reports future targets
- Reports future goals
- Reports planned improvements
- Reports intended reductions or increases
- Reports commitments to be achieved later

Example 1

Paragraph:
{pm_example}

Answer: PM


Example 2

Paragraph:
{fc_example}

Answer: FC


Paragraph to classify:

{paragraph}

Respond with EXACTLY one of:

Answer: PM

or

Answer: FC

Do not explain.
Do not provide reasoning.
"""
    return prompt


def classify_pmfc(paragraph, pm_example, fc_example):
    """Matches the reference script's classify_pmfc: generates up to 5 new
    tokens, then checks the (uppercased) output for 'PM' or 'FC'."""
    pipe, tokenizer = _get_pipe()

    prompt = generate_pmfc_prompt(paragraph, pm_example, fc_example)

    output = pipe(
        prompt,
        max_new_tokens=5,
        pad_token_id=tokenizer.eos_token_id
    )[0]["generated_text"]

    output_upper = output.strip().upper()

    if "PM" in output_upper:
        return "Answer: PM"
    if "FC" in output_upper:
        return "Answer: FC"
    return "UNKNOWN"


# =========================
# FULL PIPELINE: run on a quantitative-only subset of a DataFrame
# =========================
def classify_quantitative_paragraphs(df, qq_column):
    """
    Mirrors the reference script's full flow:
      1. Filter df to rows where qq_column == "Quantitative"
      2. Retrieve top-1 most similar PM example and top-1 most similar FC
         example (independently) for every quantitative paragraph
      3. Classify each one via one-shot prompting against Qwen2.5-1.5B
      4. Merge PMFC_Label (and the retrieved examples + similarity scores,
         needed for the per-paragraph CSV export) back into the full df

    Returns the full df with new columns added:
      PMFC_Label, EX_PM, FA_PM, PM_Similarity, EX_FC, FA_FC, FC_Similarity
    (non-quantitative rows get blank/NaN in these columns, same as the
    reference script's df["PMFC_Label"] = "" then partial assignment)
    """
    _load_pmfc_corpus()
    embedder = _get_embedder()

    df = df.copy()
    # Use object dtype with pd.NA placeholders (not "") so that the
    # numeric columns (PM_Similarity, FC_Similarity) can later receive
    # float values without hitting pandas' Arrow-backed string dtype
    # restriction (TypeError: Invalid value for dtype 'str').
    for col in ["PMFC_Label", "EX_PM", "FA_PM", "PM_Similarity", "EX_FC", "FA_FC", "FC_Similarity"]:
        df[col] = pd.Series([pd.NA] * len(df), index=df.index, dtype="object")

    quant_df = df[df[qq_column] == "Quantitative"].copy()

    if quant_df.empty:
        return df

    quant_embeddings = embedder.encode(
        quant_df["Paragraph_content"].tolist(),
        convert_to_tensor=True,
        normalize_embeddings=True
    )

    pm_scores = util.cos_sim(quant_embeddings, _CACHE["pm_embeddings"])
    pm_best_idx = torch.argmax(pm_scores, dim=1)

    fc_scores = util.cos_sim(quant_embeddings, _CACHE["fc_embeddings"])
    fc_best_idx = torch.argmax(fc_scores, dim=1)

    pm_examples = _CACHE["pm_examples"]
    fc_examples = _CACHE["fc_examples"]

    predictions = []
    ex_pm_list, ex_fc_list = [], []
    pm_sim_list, fc_sim_list = [], []

    for i in range(len(quant_df)):
        pm_idx = pm_best_idx[i].item()
        fc_idx = fc_best_idx[i].item()

        pm_text = pm_examples.iloc[pm_idx]["Paragraph_content"]
        fc_text = fc_examples.iloc[fc_idx]["Paragraph_content"]

        ex_pm_list.append(pm_text)
        ex_fc_list.append(fc_text)
        pm_sim_list.append(float(pm_scores[i, pm_idx].item()))
        fc_sim_list.append(float(fc_scores[i, fc_idx].item()))

        pred = classify_pmfc(quant_df.iloc[i]["Paragraph_content"], pm_text, fc_text)
        predictions.append(pred)

    quant_df["PMFC_Label"] = predictions
    quant_df["EX_PM"] = ex_pm_list
    quant_df["FA_PM"] = "PM"
    quant_df["PM_Similarity"] = pm_sim_list
    quant_df["EX_FC"] = ex_fc_list
    quant_df["FA_FC"] = "FC"
    quant_df["FC_Similarity"] = fc_sim_list

    for col in ["PMFC_Label", "EX_PM", "FA_PM", "PM_Similarity", "EX_FC", "FA_FC", "FC_Similarity"]:
        df.loc[quant_df.index, col] = quant_df[col]

    return df