from langchain.llms import OpenAI
from langchain import PromptTemplate, FewShotPromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector
from langchain.chains import LLMChain, SimpleSequentialChain
from langchain.prompts import load_prompt
import pandas as pd
import numpy as np
import os
import configparser

# # Get the local API Key
# config = configparser.ConfigParser()
# config.read('/Users/jmccarthy/Dev/irnd_llms/lc-agents/secrets.config')
# os.environ["OPENAI_API_KEY"] = config.get("OPENAI","key")

# # Get OpenAI instance
# llm = OpenAI(temperature=0, model_name='gpt-3.5-turbo')


def create_config(
    topic_column="topic", target_column="text", resample_rounds=1, resample_length=3
) -> dict:
    config = {
        "topic_column": topic_column,
        "target_column": target_column,
        "resample_rounds": resample_rounds,
        "resample_length": resample_length,
    }
    return config


def resample_df(
    curr_df,
    llm,
    config: dict,
    text_col: str,
):
    # default topic, title_abstract
    topics_to_test = pd.Series(curr_df[config["topic_column"]].unique())
    total_source_count = config["resample_rounds"] * config["resample_length"]
    frame = {config["topic_column"]: topics_to_test}
    topic_df = pd.DataFrame(frame)
    col_prefix = "rs"  # TODO parameterize this

    rsColNames = {}
    for i in range(0, config["resample_rounds"]):
        rsColNames[i] = f"{col_prefix}{i}"

    for tp in topics_to_test:
        matching_docs = curr_df.loc[(curr_df[config["topic_column"]] == tp)]
        if len(matching_docs) > total_source_count:
            samp_size = total_source_count
        else:
            samp_size = len(matching_docs)
        # This prevents resampling of the same document
        # TODO - handle where doc size is less than N not just less than K
        # TODO - add probability logic here --filter dataframe to probs we can use (or should this be done prior?)
        samples = np.array_split(
            matching_docs.sample(samp_size)[text_col], config["resample_length"]
        )
        source_splitter = " --*-- "
        for i in range(0, config["resample_rounds"]):
            topic_df.loc[
                topic_df[config["topic_column"]] == tp, rsColNames[i]
            ] = source_splitter.join(samples[i])

    # All resample source columns have been created/populated - Generate new labels
    # TODO - parameterize the prompt and tokenization
    resample_prompt = load_prompt("naming/prompts/resample_prompt.json")
    for i in rsColNames:
        rspPrompts = [
            resample_prompt.format(text=t, domain="science")
            for t in list(topic_df[rsColNames[i]])
        ]
        rsp_titles = [llm(prompt) for prompt in rspPrompts]
        # rsp_titles_x = [l[0].text for l in list(rsp_titles)[0][1]]
        # rsp_titles_x =  [l[2:] for l in rsp_titles_x]
        topic_df[f"{rsColNames[i]}_result"] = rsp_titles

    # Create a consenus label
    consensus_prompt = load_prompt("naming/prompts/consensus.json")
    rsc_res_name = [r + "_result" for r in rsColNames.values()]
    topic_df["consensus_source"] = topic_df[rsc_res_name].apply(
        lambda x: " ".join("[" + x + "]"), axis=1
    )
    consensus_prompts = [
        consensus_prompt.format(text=t) for t in list(topic_df["consensus_source"])
    ]
    rres = [llm(p) for p in consensus_prompts]
    # ff = [l[0].text for l in list(rres)[0][1]]
    # ff =  [l[2:] for l in ff]
    topic_df["consensus_label"] = rres

    # return our enhanced DF
    return topic_df


def create_probabilityDf(
    input_df: pd.DataFrame,
    config: dict,
    probability_threshold: float,
    id_col: str,
):
    # Set initial volume limit
    volume_limit = config["resample_rounds"] * config["resample_length"]

    # Filter based strictly on volume limit
    count_df = input_df.copy().groupby("topic")[id_col].count().reset_index()

    keep_all_topics = []
    below_limit_topics = list(
        count_df[count_df[id_col] < volume_limit]["topic"].unique()
    )
    keep_all_topics = keep_all_topics + below_limit_topics

    # Filter based on probability threshold and volume limits
    filter_prob_df = input_df[input_df["probability"] >= probability_threshold]
    count_df_probFilter = (
        filter_prob_df.copy().groupby("topic")[id_col].count().reset_index()
    )
    below_limit_topics = list(
        count_df_probFilter[count_df_probFilter[id_col] < volume_limit][
            "topic"
        ].unique()
    )
    keep_all_topics = keep_all_topics + below_limit_topics

    all_docs_df = input_df[input_df["topic"].isin(keep_all_topics)]
    return pd.concat([all_docs_df, filter_prob_df])
