import gradio as gr
import os
from transformers import AutoTokenizer
from get_loss.get_loss_hf import run_get_loss
import pdb
from types import SimpleNamespace
import pandas as pd
import plotly.express as px
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime
import scipy
<<<<<<< HEAD
=======
import shutil
>>>>>>> 95f252d5216b18bcd503dc7425ceb1dd1ccb8f6b
# os.system('git clone https://github.com/EleutherAI/lm-evaluation-harness')
# os.system('cd lm-evaluation-harness')
# os.system('pip install -e .')
# -i https://pypi.tuna.tsinghua.edu.cn/simple
# 第一个功能：基于输入文本和对应的损失值对文本进行着色展示


csv_file_path = 'data.csv'

def save_and_share_csv():
    src_path = './data/0309_merge_gjo.csv'
    dest_dir = './save/'
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)
    dest_path = os.path.join(dest_dir, '0309_merge_gjo_shared.csv')
    shutil.copy(src_path, dest_path)
    return """
    <script>
    alert('Data shared successfully! CSV saved to ./save/ directory.');
    </script>
    """
    #弹窗没有但反正能保存


def plot_ppl():
    df = pd.read_csv(csv_file_path)
    # 假设df已经有适当的列用于绘图
    fig = px.line(df, x='date', y='loss_mean_at_1000', color='model', title='PPL with Time')
    return fig


def color_text(text_list=["hi", "FreshEval","!"], loss_list=[0.1,0.7]):
    """
    根据损失值为文本着色。
    """
    highlighted_text = []
    # print('loss_list',loss_list)
    # ndarray to list
    loss_list = loss_list.tolist()
    loss_list=[0]+loss_list
    # print('loss_list',loss_list)
    # print('text_list',text_list)
    # pdb.set_trace()
    for text, loss in zip(text_list, loss_list):
        # color = "#FF0000" if float(loss) > 0.5 else "#00FF00"
        color=loss/20#TODO rescale
        # highlighted_text.append({"text": text, "bg_color": color})
        highlighted_text.append((text, color))
        print('highlighted_text',highlighted_text)
    return highlighted_text

# 第二个功能：根据 ID 列表和 tokenizer 将 ID 转换为文本，并展示
def get_text(ids_list=[0.1,0.7], tokenizer=None):
    """
    给定一个 ID 列表和 tokenizer 名称，将这些 ID 转换成文本。
    """
    # return ['Hi', 'Adam']
    # tokenizer = AutoTokenizer.from_pretrained(tokenizer)
    # print('ids_list',ids_list)
    # pdb.set_trace()
    text=[]
    for id in ids_list:
        text.append(  tokenizer.decode(id, skip_special_tokens=True))
    # 这里只是简单地返回文本，但是可以根据实际需求添加颜色或其他样式
    print(f'L41:{text}')
    return text


# def get_ids_loss(text, tokenizer, model):
#     """
#     给定一个文本，model and its tokenizer,返回其对应的 IDs 和损失值。
#     """
#     # tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
#     # model = AutoModelForCausalLM.from_pretrained(model_name)
#     # 这里只是简单地返回 IDs 和损失值，但是可以根据实际需求添加颜色或其他样式
#     return [1, 2], [0.1, 0.7]

def harness_eval(question,  answer_index, answer_type, model=None,*choices,):
    '''
    use harness to test one question,  can specify the model, (extract or ppl)
    '''
    # print(f'question,choices,answer_index,model,tokenizer: {question,choices,answer_index,model,tokenizer}')
    print(f'type of choices: {type(choices)} and type of choices[0]: {type(choices[0])}')
    print(f'choices: {choices}')
    # TODO add the model and its score
    # torch.nn.functional.softmax(output.logits, dim=0)
    # topk = torch.topk(output.logits, 5)
    
    return {'A':0.5, 'B':0.3, 'C':0.1, 'D':0.1}




def plotly_plot_text():#(df, x, y, color,title, x_title, y_title):
        # plotly_plot(sample_df, 'date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl')
    df=pd.read_csv('./data/tmp.csv')
    df['date'] = pd.to_datetime(df['date'])
    # sort by date
    df.sort_values(by='date', inplace=True)

    # use a dic to filter the dataframe
    df = df[df['file_name'] == 'arxiv_computer_science']

    x,y,color,title, x_title, y_title='date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl'

    fig = px.line(df, x=x, y=y, color=color,title=title)
    fig.update_xaxes(title_text=x_title)
    fig.update_yaxes(title_text=y_title)
    # fig.update_layout()
    return fig

def plotly_plot_question(use_start=True,gjo=True,time_diff=False):#(df, x, y, color,title, x_title, y_title):
        # plotly_plot(sample_df, 'date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl')
    data=pd.read_csv('./data/0309_merge_gjo.csv')
# Model_x,Release Date,model,MMLU,GSM8,Humanities,SocialSciences,STEM,Other,Longbench,Question,Model_y,Start Time,End Time,Acc,Right Possibility

    data['date'] = pd.to_datetime(data['End Time'])
    # sort by date
    data.sort_values(by='date', inplace=True)

    # use a dic to filter the dataframe
    # df = df[df['file_name'] == 'arxiv_computer_science']

    # x,y,color,title, x_title, y_title='date', 'Right Possibility', 'model','Right Possibility with time', 'time', 'Right Possibility'

    # fig = px.line(df, x=x, y=y, color=color,title=title)
    # fig.update_xaxes(title_text=x_title)
    # fig.update_yaxes(title_text=y_title)
    if not use_start:
        data['Start Time']=data['End Time']

    # # Convert the 'Release Date' and 'Start Time' columns to datetime
    data['Release Date'] = pd.to_datetime(data['Release Date'])
    data['Start Time'] = pd.to_datetime(data['Start Time'])

    data_cleaned = data.dropna(subset=['Release Date', 'Start Time'])
    if time_diff:
        if gjo:
            data_cleaned['Time Difference (Months)'] = ((data_cleaned['Start Time'] - data_cleaned['Release Date']) / pd.Timedelta(days=90)).round().astype(int)
        else:
            data_cleaned['Time Difference (Months)'] = ((data_cleaned['Start Time'] - data_cleaned['Release Date']) / pd.Timedelta(days=365)).round().astype(int)
    else:
        time_point= datetime(2015, 1, 1)
        data_cleaned['Time Difference (Months)'] = ((data_cleaned['Start Time'] - time_point) / pd.Timedelta(days=90)).round().astype(int)
    # Step 1: Fill missing months with linear interpolation (if necessary)
    # Note: This dataset might not have explicit missing months, but we will ensure continuity for plotting
    # pdb.set_trace()
    # data_cleaned
    # data_cleaned['Time Difference (Months)'].value_counts()
    # Ensure 'Time Difference (Months)' is sorted for each model before applying rolling mean
    data_cleaned.sort_values(by=['Model_x', 'Time Difference (Months)'], inplace=True)

    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    import plotly.express as px
    from scipy.interpolate import CubicSpline


    # Initialize figure with subplots
    # fig = make_subplots(rows=2, cols=1, subplot_titles=('Accuracy (Acc)', 'Right Possibility'))
    # make this pic large enough
    fig = make_subplots(rows=2, cols=1, subplot_titles=('Accuracy (Acc)', 'Right Possibility'),vertical_spacing=0.1)


    colors = px.colors.qualitative.Plotly  # Use Plotly's qualitative colors for consistency

    # Iterate over each unique model to plot their data
    for i, (model_name, group) in enumerate(data_cleaned.groupby('Model_x')):
        color = colors[i % len(colors)]  # Cycle through colors
        #  mean accuracy and right possibility for each model
        group=group.groupby(['Model_x', 'Time Difference (Months)'])\
            .agg({'Acc':'mean','Right Possibility':'mean','Release Date':'first','Start Time':'first'}).reset_index()
        
        # Divide the data into before and after based on 'Release Date' and 'Start Time'
        before = group[group['Release Date'] >= group['Start Time']]
        after = group[group['Release Date'] < group['Start Time']]
        
        # Concat the last row of 'before' to 'after' if 'before' is not empty
        if not before.empty:
            after = pd.concat([before.iloc[[-1]], after])

        # # ================================================================================
        # before = CubicSpline(before['Time Difference (Months)'], before['Acc'])
        # after = CubicSpline(after['Time Difference (Months)'], after['Acc'])

        # before = CubicSpline(before['Time Difference (Months)'], before['Right Possibility'])
        # after = CubicSpline(after['Time Difference (Months)'], after['Right Possibility'])
        # # ================================================================================

        
        # Plot 'Acc' on the first subplot
        fig.add_trace(go.Scatter(x=before['Time Difference (Months)'], y=before['Acc'], mode='lines', name=model_name + ' (Acc before)', line=dict(color=color)), row=1, col=1)
        fig.add_trace(go.Scatter(x=after['Time Difference (Months)'], y=after['Acc'], mode='lines', name=model_name + ' (Acc after)', line=dict(color=color, dash='dash')), row=1, col=1)
        
        # Plot 'Right Possibility' on the second subplot
        fig.add_trace(go.Scatter(x=before['Time Difference (Months)'], y=before['Right Possibility'], mode='lines', name=model_name + ' (Right Possibility before)', line=dict(color=color)), row=2, col=1)
        fig.add_trace(go.Scatter(x=after['Time Difference (Months)'], y=after['Right Possibility'], mode='lines', name=model_name + ' (Right Possibility after)', line=dict(color=color, dash='dash')), row=2, col=1)

    # Update layout if needed
    fig.update_layout(height=600, width=800, title_text="Model Performance Over Time")
    # fig.update_layout()
    return fig

# def plotly_plot(df, x, y, color, title, x_title, y_title):
#     fig = px.line(df, x=x, y=y, color=color, title=title)
#     fig.update_xaxes(title_text=x_title)
#     fig.update_yaxes(title_text=y_title)
#     return fig

def show_attention_plot(model_name,texts):
    # 初始化分词器和模型，确保在模型配置中设置 output_attentions=True
    args=SimpleNamespace(texts=texts,model=model_name)
    print(f'L60,text:{texts}')
    rtn_dic=run_get_loss(args)
    # print(rtn_dic)
    # pdb.set_trace()
    # {'logit':logit,'input_ids':input_chunk,'tokenizer':tokenizer,'neg_log_prob_temp':neg_log_prob_temp}
    # ids, loss =rtn_dic['input_ids'],rtn_dic['loss']#= get_ids_loss(text, tokenizer, model)
    # notice here is numpy ndarray
    tokenizer, model = rtn_dic['tokenizer'],rtn_dic['model'] 
    text = "Here is some text to encode"

    # 使用分词器处理输入文本
    inputs = tokenizer(text, return_tensors="pt")
    # 进行前向传播，获取输出
    outputs = model(**inputs, output_attentions=True)

    # 检查是否成功获得了 attentions
    if "attentions" in outputs:
        last_layer_attentions = outputs.attentions[-1]  # 获取最后一层的 attention 矩阵
        print("Successfully retrieved the attention matrix:", last_layer_attentions.shape)
    else:
        pdb.set_trace()
        print("Attention matrix not found in outputs.")

    # 假设 last_layer_attentions 是我们从模型中提取的注意力矩阵
    # last_layer_attentions 的形状应该是 [batch_size, num_heads, seq_length, seq_length]
    # 为了简化，我们这里只查看第一个样本、第一个头的注意力矩阵
    attention_matrix = last_layer_attentions[0, 0].detach().numpy()

    # 使用 matplotlib 绘制热图
    plt.figure(figsize=(10, 8))
    plt.imshow(attention_matrix, cmap='viridis')

    # 添加标题和标签以提高可读性
    plt.title('Attention Matrix Visualization')
    plt.xlabel('Tokens in Sequence')
    plt.ylabel('Tokens in Sequence')

    # 添加颜色条
    plt.colorbar()

    # 保存图表到文件
    # plt.savefig('/223040239/medbase/attention_matrix_visualization.png')
    return plt


def color_pipeline(texts=["Hi","FreshEval","!"],  model=None):
    """
    给定一个文本，返回其对应的着色文本。
    """
    print('text,model',texts,model)
    args=SimpleNamespace(texts=texts,model=model)
    print(f'L60,text:{texts}')
    rtn_dic=run_get_loss(args)
    # print(rtn_dic)
    # pdb.set_trace()
    # {'logit':logit,'input_ids':input_chunk,'tokenizer':tokenizer,'neg_log_prob_temp':neg_log_prob_temp}
    ids, loss =rtn_dic['input_ids'],rtn_dic['loss']#= get_ids_loss(text, tokenizer, model)
    # notice here is numpy ndarray
    tokenizer=rtn_dic['tokenizer'] # get tokenizer
    text = get_text(ids, tokenizer)
    # print('ids, loss ,text',ids, loss ,text)
    return color_text(text, loss)


# TODO can this be global ? maybe need session to store info of the user
# visible_btn_num = 4

# 创建 Gradio 界面
with gr.Blocks() as demo:
    # visible_btn_num = 4
    model_input = gr.Textbox(label="model name", placeholder="input your model name here... now I am trying phi-2...")#TODO make a choice here


    with gr.Tab("color your text"):
        with gr.Row():
            text_input = gr.Textbox(label="input text", placeholder="input your text here...")
            # file_input = gr.File(file_count="multiple",label='to add content')#
            # TODO craw and drop the file

            # loss_input = gr.Number(label="loss")
            # model_input = gr.Textbox(label="model name", placeholder="input your model name here... now I am trying phi-2...")#TODO make a choice here
        output_box=gr.HighlightedText(label="colored text")#,interactive=True
        
        gr.Examples(
                [
                    ["Hi FreshEval !", "microsoft/phi-2"],
                    ["Hello FreshBench !", "/home/sribd/chenghao/models/phi-2"],
                ],
                [text_input, model_input],)
            #     cache_examples=True,
            #     # cache_examples=False,
            #     fn=color_pipeline,
            #     outputs=output_box
            # )
            # TODO select models that can be used online
            # TODO maybe add our own models


        color_text_output = gr.HTML(label="colored text")
        color_text_button = gr.Button("color the text").click(color_pipeline, inputs=[text_input, model_input], outputs=output_box)

        # markdown
        gr.Markdown('### How to use this app')


        attention_plot=gr.Plot(label='attention plot')
        see_attention_button = gr.Button("see attention").click(show_attention_plot,inputs=[model_input, text_input],outputs=[attention_plot])
        
        




        date_time_input = gr.Textbox(label="the date when the text is generated")#TODO add date time input
        description_input = gr.Textbox(label="description of the text")
        submit_button = gr.Button("submit a post or record").click()
        #TODO add model and its score

    with gr.Tab('test your qeustion'):
        '''
        use extract, or use ppl
        '''
        question=gr.Textbox(label="input question", placeholder='input your question here...')
        answer_index=gr.Textbox(label="right answer index", placeholder='index for right anser here, start with 0')#TODO add multiple choices, 
        # model_input = gr.Textbox(label="model name", placeholder="input your model name here... now I am trying phi-2...")#TODO make a choice here

        btn_list = []

        # choices=gr.Textbox(placeholder='input your other choices here...')
        
        button_limit=10
        # global visible_btn_num
        visible_btn_num = 4

        from gradio_samples.add_components import add_one_btn, remove_one_btn, get_text_content

        # use partial function
        from functools import partial
        add_one_btn=partial(add_one_btn,button_limit=button_limit,)#visible_btn_num = visible_btn_num)
        remove_one_btn=partial(remove_one_btn,button_limit=button_limit,)#visible_btn_num = visible_btn_num)
            
        # with gr.Blocks() as demo:
        with gr.Row():
            for i in range(button_limit):
                if i<visible_btn_num:
                    btn = gr.Textbox(visible=True)
                else:
                    btn = gr.Textbox(visible=False)
                btn_list.append(btn)
        b = gr.Button("add_one_choice(make sure every existing choice is filled)")
        print(f'len(btn_list): {len(btn_list)}')
        b.click(add_one_btn, btn_list, btn_list)
        b = gr.Button("remove_one_choice")
        b.click(remove_one_btn, btn_list, btn_list)

        # # print(f'len(btn_list): {len(btn_list)}')

        # print('btn_list is ',type(btn_list),btn_list)


        # b = gr.Button("Get Text Content")
        # output = gr.Textbox()
        # b.click(get_text_content, btn_list, output)

        
        # test_button=gr.Button('test').click(harness_eval())# TODO figure out the input and output

        answer_type=gr.Dropdown(label="answer type", choices=['extract', 'ppl'])
        #TODO add the model and its score
        answer_label=gr.Label('the answers\'s detail')# RETURN the answer and its score,in the form of dic{str: float}

        test_question_button=gr.Button('test question').click(harness_eval,inputs=[question, answer_index ,model_input,answer_type,*btn_list],outputs=[answer_label])

        forecast_q='A Ukrainian counteroffensive began in 2023, though territorial gains by November 2023 were limited (Economist, BBC, Newsweek). The question will be suspended on 31 July 2024 and the outcome determined using data as reported in the Brookings Institution\'s "Ukraine Index" (Brookings Institution - Ukraine Index, see "Percentage of Ukraine held by Russia" chart). If there is a discrepancy between the chart data and the downloaded data (see "Get the data" within the "NET TERRITORIAL GAINS" chart border), the downloaded data will be used for resolution.'
        # answer_list=['Less than 5%','At least 5%, but less than 10%','At least 10%, but less than 15%','At least 15%, but less than 20%','20% or more' ]
        answer_list=['Less than 5%','At least 5%, but less than 10%','At least 10%, but less than 15%','15% or more' ]
         
        # gr.Examples([
        #             [forecast_q, '&&&&&&'.join(answer_list), '0']
        #              ],
        #             [question, choices, answer_index])
        gr.Examples([
            [forecast_q, answer_list[0],answer_list[1],answer_list[2],answer_list[3], '0']
                ],
            [question,btn_list[0],btn_list[1],btn_list[2],btn_list[3], answer_index])

        date_time_input = gr.Textbox(label="the date when the text is generated")#TODO add date time input
        description_input = gr.Textbox(label="description of the text")
        submit_button = gr.Button("submit a post or record").click()

        #TODO add the model and its score

        def test_question(question, answer, other_choices):
            '''
            use extract, or use ppl
            '''
            answer_ppl, other_choices_ppl =  (question, answer, other_choices)
            return answer_ppl, other_choices_ppl



    with gr.Tab("model text ppl with time"):
        '''
        see the matplotlib example, to see ppl with time, select the models
        '''
        # load the json file with time,

        # sample_df=pd.DataFrame({'time':pd.date_range('2021-01-01', periods=6), 'ppl': [1,2,3,4,5,6]})
        # pd_df=pd.read_csv('./data/tmp.csv')
        # pd_df['date'] = pd.to_datetime(pd_df['date'])
        # print(pd_df.head)
        # # gr_df=gr.Dataframe(pd_df)
        # gr_df=pd_df


        # print(gr_df.head)
        # print('done')
        # sample
        plot=gr.Plot(label='model text ppl')
        # plotly_plot(gr_df, 'date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl')
        # draw_pic_button=gr.Button('draw the pic').click(plotly_plot,inputs=['gr_df', 'date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl'],outputs=[plot])
        draw_pic_button=gr.Button('draw the pic').click(plotly_plot_text,inputs=[],outputs=[plot])


    
    with gr.Tab("model quesion acc with time"):
        '''
        see the matplotlib example, to see ppl with time, select the models
        ''' 

        # pd_df=pd.read_csv('./data/meta_gjo_df.csv')
        # pd_df['date'] = pd.to_datetime(pd_df['end_date'])
        # print(pd_df.head)
        # gr_df=gr.Dataframe(pd_df)
        # gr_df=pd_df


        # print(gr_df.head)
        # print('done')
        # sample
        plot=gr.Plot(label='question acc with time')
        # plotly_plot(gr_df, 'date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl')
        # draw_pic_button=gr.Button('draw the pic').click(plotly_plot,inputs=['gr_df', 'date', 'loss_mean_at_1000', 'model','ppl with time', 'time', 'ppl'],outputs=[plot])
        draw_pic_button=gr.Button('draw the pic').click(plotly_plot_question,inputs=[],outputs=[plot])




    with gr.Tab("hot questions"):
        '''
        see the questions and answers
        '''
        with gr.Tab("ppl"):
            '''
            see the questions
            '''
        
<<<<<<< HEAD

demo.launch(debug=True)

=======
    with gr.Row():
        plot_btn = gr.Button("Generate Plot")
        share_btn = gr.Button("Share Data")
    with gr.Row():
        plot_space = gr.Plot()
        share_result = gr.Textbox(visible=False)
    
    # 当点击“Generate Plot”按钮时，调用plotly_plot_question函数并在plot_space显示结果
    plot_btn.click(fn=plotly_plot_question, inputs=[], outputs=plot_space)
    
    # 当点击“Share Data”按钮时，调用save_and_share_csv函数并在share_result显示结果
    share_btn.click(fn=save_and_share_csv, inputs=[], outputs=share_result)


demo.launch(share=True,debug=True)
>>>>>>> 95f252d5216b18bcd503dc7425ceb1dd1ccb8f6b