#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Aug 15 01:23:38 2025

@author: samd
"""
import bilbystats as bs
import numpy as np
import pandas as pd
from typing import List, Tuple, Optional
import pickle 

async def optimize_one_shot(train_df, instructions, mymessages = [], model_name = 'gpt-4o-mini'):
    ntrain = len(train_df)
    accuracy = []
    
    for i in range(ntrain):
        sub_df = train_df.drop(index=i)

        oneshotmessage = mymessages + bs.combine_messages(train_df['text'][i], train_df['label'][i], user_prefix = 'Text: ', chat_prefix = 'Label: ')
        run = await bs.batch_llm_api_async(sub_df.tolist(),
                                    instructions = instructions,
                                    model_name = 'gpt-4o-mini',
                                    temperature = 0, prefix = 'Text: ', 
                                    mymessages = mymessages)
        
        run['num_label'] = run['label'].str.replace('Label: ', '', regex=False).astype(int)
        accuracy_df = bs.accuracy_df(train_df['label'], run['num_label'])
        accuracy_df.index = accuracy_df['Label']
        accuracy.append = accuracy_df['Accuracy']['Overall']
        
    min_i = int(np.argmin(accuracy))
    mymessages = mymessages + bs.combine_messages(train_df['text'][i], train_df['label'][i], user_prefix = 'Text: ', chat_prefix = 'Label: ')
    
    return mymessages, min_i

async def optimize_few_shot(train_df, instructions, mymessages = [], model_name = 'gpt-4o-mini', nshots = 5):
    for i in range(nshots):
        mymessages, min_i = optimize_one_shot(train_df, instructions, mymessages, model_name)
        sub_df = train_df.drop(index=min_i)
        
    return mymessages

def nshot_messages(df, labels, label_names, text_covariate, user_prefix = 'Text: ', chat_prefix = 'Output: \n', type2use = 'message'):
    if type2use == 'str':
        messages = ''
    elif type2use == 'message':
        messages = []
    else:
        messages = []
       
    store_texts = []
    store_cats = []
    for i in range(len(df)): 
        labelled_stuff = []
        for j in range(len(labels)):
            labelled_stuff.append(df[labels[j]].iloc[i])
            
        cat2use = bs.cat2output(label_names, labelled_stuff)
        store_cats.append(cat2use)
        store_texts.append(df[text_covariate].iloc[i])
        
    messages = list2nshot(store_texts, store_cats, user_prefix = user_prefix, chat_prefix = chat_prefix, type2use = type2use)
        
    return messages

def set_of_nshots(df, labels, label_names, text_covariate, user_prefix = 'Text: ', type2use = 'message'):
    set_of_nshot_messages = []
    for i in range(len(df)):
        df_subset = df.iloc[:i].reset_index(drop=True)
        messages = nshot_messages(df_subset, labels, label_names, text_covariate, user_prefix, type2use)
        set_of_nshot_messages.append(messages)
            
    return set_of_nshot_messages

def list2nshot(texts, labels, user_prefix = 'Text: ', chat_prefix = '', type2use = 'message'):
    if type2use == 'str':
        messages = ''
    elif type2use == 'message':
        messages = []
    else:
        messages = []
        
    if type2use == 'message':
        for i in range(len(texts)):
            messages += bs.combine_messages([texts[i]], [str(labels[i])], user_prefix = user_prefix)
    elif type2use == 'str':         
        for i in range(len(texts)):
            messages += user_prefix + texts[i] + '\n'
            messages += chat_prefix + str(labels[i]) + '\n'
    else:
        for i in range(len(texts)):
            messages += [user_prefix +  texts[i], chat_prefix + str(labels[i])]
            
    if type2use == 'str':
        messages = messages[:-1]
            
    return messages

async def linearbs(train_df, valid_df, instructions, output_format, labels, 
                   label_names, text_covariate = 'text', label_covariate = 'label', model_name = 'gpt-4o-mini', 
                    user_prefix = 'Text: ', chat_prefix = 'Output: \n', type2use = 'message'):
    
    # Compute the original accuracy
    _, run_accuracy = await bs.test_prompts(
        data = valid_df,
        list_of_prompts = [instructions], 
        model_name = model_name, 
        output_format = output_format,
        prefix = user_prefix,
        text_covariate = text_covariate,
        label = label_covariate,
        accuracy_df = False
    ) 
    
    current_accuracy = run_accuracy[0]
    accuracy_store = [current_accuracy]
    succcessful_accuracy = [current_accuracy]
    current_messages = []
    
    for i in range(len(train_df)):
        print(i)
        #bs.loader(i, len(train_df))
        dfrow = train_df.iloc[i:(i+1)].reset_index(drop=True)
        new_messages = bs.nshot_messages(dfrow, labels, label_names, text_covariate, user_prefix, chat_prefix, type2use)
            
        if type2use == 'message':
            mymessages = current_messages + new_messages
            prompt2use = instructions
        else:
            prompt2use = instructions + '''Here are some examples
            ''' + current_messages + '\n' + new_messages
            mymessages = None
            
        _, run_accuracy = await bs.test_prompts(
            data = valid_df, 
            list_of_prompts = [prompt2use], 
            model_name = model_name, 
            output_format = output_format,
            prefix = user_prefix,
            text_covariate = text_covariate,
            label = label_covariate,
            mymessages = mymessages,
            accuracy_df = False
        ) 
        
        new_accuracy = run_accuracy[0]
        accuracy_store.append(new_accuracy)
            
        if new_accuracy > current_accuracy:
            if type2use == 'message':
                current_messages += new_messages
            else:
                current_messages += '\n' + new_messages
                
            succcessful_accuracy.append(new_accuracy)
            current_accuracy = new_accuracy

        out = {
            'messages': current_messages, 
            'accuracy_store': accuracy_store,
            'succcessful_accuracy': succcessful_accuracy
        }
    
    return out
    
async def my_few_shot_success(train_df, instructions, nshots = None, other_labels = [], text_covariate = 'text', label_covariate = 'label',
                           model_name = 'gpt-4o-mini', temperature = 0,
                           index_type = 'first', ones_that_work = True):
    messages = []
    
    if not nshots:
        nshots = len(train_df)
    
    df = await bs.batch_llm_api_async_mr(train_df[text_covariate].tolist(),
                                    instructions = instructions,
                                    model_name = model_name,
                                    temperature = temperature,
                                    prefix = 'Text: \n')
    
    matching_labels = (train_df[label_covariate] == df[label_covariate])
    
    set_of_labels = []
    other_labels.append([label_covariate])
    for i in range(len(train_df)):
        if type(matching_labels[i]) != pd._libs.missing.NAType:
            if matching_labels[i] == ones_that_work:
                text = [train_df[text_covariate].iloc[i]]
                label = [train_df[label_covariate].iloc[i]]
                set_of_labels.append(label[0])
                
                labelled_stuff = []
                for j in range(other_labels):
                    labelled_stuff.append(train_df[other_labels[j]].iloc[i])
                cat2use = bs.cat2output(other_labels, labelled_stuff)
                                        
                one_shot_message = bs.combine_messages(text, cat2use, user_prefix = 'Text: \n')
                messages += one_shot_message
            
    if index_type == 'first':
        # Restrict to the first n shots
        returned_messages = messages[:2*nshots]
    else:
        unique_labels = list(set(set_of_labels))
        
        label_indices = []
        for item in unique_labels:
           indices = [i for i, x in enumerate(set_of_labels) if x == item][:nshots]
           label_indices += indices
          
        output_indices = []
        for idx in label_indices:
            output_indices.append(2*idx)
            output_indices.append(2*idx + 1)
            
        returned_messages = [messages[i] for i in output_indices]
        
    return returned_messages

async def few_shot_success(train_df, instructions, nshots = None, 
                           model_name = 'gpt-4o-mini', temperature = 0,
                           index_type = 'first', ones_that_work = True):
    messages = []
    
    if not nshots:
        nshots = len(train_df)
    
    run = await bs.batch_llm_api_async(train_df['text'].tolist(),
                                    instructions = instructions,
                                    model_name = model_name,
                                    temperature = temperature,
                                    prefix = '[[ ## text ## ]]\n')
   
    df = pd.DataFrame(run, columns=["raw"])
    df["label"] = df["raw"].str.extract(r'\n(\d)\n')
    df["label"] = df["label"].astype("Int64")
    
    matching_labels = (train_df['label'] == df['label'])
    
    set_of_labels = []
    for i in range(len(train_df)):
        if type(matching_labels[i]) != pd._libs.missing.NAType:
            if matching_labels[i] == ones_that_work:
                text = [train_df["text"].iloc[i]]
                label = [train_df["label"].iloc[i]]
                set_of_labels.append(label[0])
                one_shot_message = bs.combine_messages(text, label, user_prefix = '[[ ## text ## ]]\n', chat_prefix = '[[ ## label ## ]]\n', user_suffix = '\n', chat_suffix = '\n[[ ## completed ## ]]')
                messages += one_shot_message
            
    if index_type == 'first':
        # Restrict to the first n shots
        returned_messages = messages[:2*nshots]
    else:
        unique_labels = list(set(set_of_labels))
        
        label_indices = []
        for item in unique_labels:
           indices = [i for i, x in enumerate(set_of_labels) if x == item][:nshots]
           label_indices += indices
          
        output_indices = []
        for idx in label_indices:
            output_indices.append(2*idx)
            output_indices.append(2*idx + 1)
            
        returned_messages = [messages[i] for i in output_indices]
        
    return returned_messages

async def bootstrap_few_shot_dep(train_df, instructions, nboot = 5, model_name = 'gpt-4o-mini'):
    set_of_messages = []
    for i in range(nboot):
        text = [train_df["text"].iloc[i]]
        label = [train_df["label"].iloc[i]]
        one_shot_message = bs.combine_messages(text, label, user_prefix = '[[ ## text ## ]]\n', chat_prefix = '[[ ## label ## ]]\n', user_suffix = '\n', chat_suffix = '\n[[ ## completed ## ]]')
        set_of_messages.append(one_shot_message)
    
def nshot_messages_orig(data_df, n_vec, covariate = 'text', label = 'label', setting = 'json',  user_prefix = 'Text: '):
    
    set_of_messages = []
    
    if setting == 'dspy':
        for i in range(len(n_vec)):
            n = n_vec[i]
            messages = bs.combine_messages(data_df[covariate].iloc[:n].tolist(), data_df[label].iloc[:n].tolist(), user_prefix = '[[ ## text ## ]]\n', chat_prefix = '[[ ## label ## ]]\n', user_suffix = '\n', chat_suffix = '\n[[ ## completed ## ]]')
            set_of_messages.append(messages)
    elif setting == 'json':
        for i in range(len(n_vec)):
            n = n_vec[i]
            list_of_labels = [ [x] for x in data_df[label].iloc[:n].tolist() ]
            chat_messages = bs.multiple_cats(['label'], list_of_labels)
            messages = bs.combine_messages(data_df[covariate].iloc[:n].tolist(), chat_messages)
            set_of_messages.append(messages)
        
    return set_of_messages

async def test_messages(
    data: pd.DataFrame, 
    instructions,
    list_of_messages: List[str], 
    model_name: str, 
    covariate = 'text',
    saveloc = './save.pkl',
    prefix: str = '',
) -> Tuple[pd.DataFrame, List[pd.DataFrame]]:
    """
    Test multiple prompts on the same dataset and return predictions and accuracy metrics.

    This function evaluates a list of prompts by running them on the provided data
    using the specified model, then calculates accuracy metrics for each prompt.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data : pd.DataFrame
        DataFrame containing test data with 'text' and 'label' columns.
    list_of_prompts : List[str]
        List of prompts to test on the data.
    model_name : str
        Name of the model to use for generating predictions.
    output_format : str, optional
        Expected output format specification to append to each prompt (default is '').

    ---------------------------------------------------------------------------
    OUTPUT:
    run_df : pd.DataFrame
        DataFrame with columns 'prompt_0', 'prompt_1', etc., containing predictions
        from each prompt.
    run_accuracy : List[pd.DataFrame]
        List of accuracy DataFrames (one per prompt) containing detailed metrics
        including accuracy, precision, recall, and F1-score for each label.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """

    full_data = []
    accuracy = []
    precision = []
    recall = []
    f1_score = []
    prompt_length = []
    
    for i in range(len(list_of_messages)):
        message = list_of_messages[i]
        acc = await bs.eval_im(instructions, message, data, model_name = model_name, covariate = covariate)
        acc.index = acc['Label']
        accuracy.append(acc['Accuracy']['Overall'])
        precision.append(acc['Precision']['Overall'])	
        recall.append(acc['Recall']['Overall'])
        f1_score.append(acc['F1_Score']['Overall'])	
        full_data.append(acc)
        pl = 0
        for im in message:
            pl += len(im)
            
        prompt_length.append(pl)
            
    session_data = {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1_score": f1_score,
        "full_data": full_data,
        "list_of_messages": list_of_messages,
        "instructions": instructions
    }

    # Save to file
    with open(saveloc, "wb") as f:
        pickle.dump(session_data, f)
        
    bs.plot_metrics(accuracy, precision, recall, f1_score, pl)

    return accuracy, precision
    