import numpy as np
import pandas as pd
from typing import Iterable, Dict, List

from src.llm import generate
from src.retriever import retrieve
from src.prompts import build_vc_analyzer_prompt
from src.query_engine import _parse_vc_analyzer_result
from .investments import normalize_scores

def score_value_chains(vc_dist_ref:pd.DataFrame,
                        val_chain_col:str,
                        ruling_opp_swing_share:Iterable):
    vc_dist_ref["score"] = 0
    vc_dist_ref.loc[(vc_dist_ref[val_chain_col].isin(ruling_opp_swing_share))&(vc_dist_ref['category']=="NRM safe")&(vc_dist_ref['share']>0.5), "score"] = 4
    vc_dist_ref.loc[~vc_dist_ref[val_chain_col].isin(ruling_opp_swing_share), 'score'] = 3
    vc_dist_ref.loc[(vc_dist_ref[val_chain_col].isin(ruling_opp_swing_share))&(vc_dist_ref['category']=="Swing")&(vc_dist_ref['share']>0.5), "score"] = 2
    vc_dist_ref.loc[(vc_dist_ref[val_chain_col].isin(ruling_opp_swing_share))&(vc_dist_ref['category']=="Opp. safe")&(vc_dist_ref['share']>0.5), "score"] = 1

    return vc_dist_ref[vc_dist_ref['score']>0]

def categorize_vc_production(vc_dist_ref:pd.DataFrame,
                             val_chain_col:str):
    vc_dist_ref = vc_dist_ref[~vc_dist_ref['category'].str.lower().str.contains("plurality")].reset_index(drop=True)
    ros_vc = vc_dist_ref.groupby([val_chain_col]).apply(lambda x: x[x["share"]>0.5]).reset_index()[val_chain_col].unique()

    score_df = score_value_chains(vc_dist_ref, val_chain_col, ros_vc)
    score_df = score_df.groupby(['score'], as_index=False)[val_chain_col].unique().explode(val_chain_col).reset_index(drop=True)

    score_df = score_df[[val_chain_col, 'score']]
    return normalize_scores(score_df)

def analyze_committment_to_vc(cfg:Dict, value_chain:str):
    search_query = f"{value_chain} production in uganda"
    retrieval_result = retrieve(search_query, cfg, route="both")
    context = retrieval_result.get("context", "")

    prompt = build_vc_analyzer_prompt(cfg.get("policy_docs", []), context, value_chain, cfg["chatbot"].get("domain"))
    raw = generate(
        "",
        prompt,
        cfg,
        max_tokens=4096,
        )

    return _parse_vc_analyzer_result(raw)

def analyze_budget_trend(data:pd.DataFrame):
    df = data.copy()
    df["trend"] = (df["financial_commitment.budget_end_value"].replace("", np.nan).astype(float)-df["financial_commitment.budget_start_value"].replace("", np.nan).astype(float))/df["financial_commitment.budget_start_value"].replace("", np.nan).astype(float)
    conditions = [
        (df["trend"] <= -0.5),
        (df["trend"] > -0.5) & (df["trend"] <= -0.1),
        (df["trend"] > -0.1) & (df["trend"] <= 0.1),
        (df["trend"] > 0.1) & (df["trend"] <= 0.5),
        (df["trend"] > 0.5),
    ]

    # 2. Define the corresponding choices/mapped values
    choices = range(1, len(conditions)+1, 1)

    # 3. Apply numpy.select with a default fallback value
    df["trend"] = np.select(conditions, choices, default=np.nan)
    return df, conditions

def get_scores(data:pd.DataFrame, conditions:List):
    df = data.copy()
    results = df[["value_chain", "trend"]].reset_index(drop=True)

    strat_imp_cols = ["strategic_commitment.development_plan.value", "strategic_commitment.agricultural_strategy.value", 
                    "strategic_commitment.flagship_project.value", "strategic_commitment.standalone_vc_strategy.value"]

    results["strategic_importance"] = df[strat_imp_cols].mean(axis=1).reset_index(drop=True)
    results["depth"] = df["strategic_commitment.documents_discussing_value_chain"]-df[strat_imp_cols].sum(axis=1).reset_index(drop=True)
    results["breadth"] = (df["strategic_commitment.documents_discussing_value_chain"]/df["strategic_commitment.documents_examined"]).reset_index(drop=True)

    results["metric"] = (0.5*results["breadth"])+(0.5*results["strategic_importance"]).reset_index(drop=True)
    results["coverage"] = (results["depth"]/(df["strategic_commitment.documents_examined"]-df[strat_imp_cols].sum(axis=1))).reset_index(drop=True)
    results["strategic_importance_score"] = (0.8*df[strat_imp_cols].sum(axis=1))+(0.2*results["breadth"]).reset_index(drop=True)
    results["institutional_committment_score"] = df[[c for c in df.columns if "value" in c and "institutional_commitment" in c]].sum(axis=1)

    return normalize_committment_scores(results, conditions)

def normalize_committment_scores(results:pd.DataFrame, conditions:List):
    
    results["trend_normalized"] = (results["trend"]-min(1, results["trend"].min()))/max(len(conditions), results["trend"].max())
    results["strategic_importance_normalized"] = results["strategic_importance_score"]/3.4
    results["institutional_commitment_normalized"] = results["institutional_committment_score"]/3

    return results

def analyze_vc_committment(llm_outputs:Dict):
    print(llm_outputs)
    c2vc_df = pd.json_normalize(llm_outputs.values())
    c2vc_df.insert(0, "value_chain", llm_outputs.keys())
    
    c2vc_df = c2vc_df.replace("YES", 1)
    c2vc_df = c2vc_df.replace("NOT_FOUND", 0)
    c2vc_df = c2vc_df.replace("NO", 0)

    trend, conditions = analyze_budget_trend(c2vc_df)
    return get_scores(trend, conditions)