import os 
import re 
import matplotlib .pyplot as plt 
import matplotlib .ticker as mticker 

def parse_accuracy_data (directory_path ):
    """
    指定されたディレクトリから精度のデータを抽出します。
    ファイル名からイテレーション数を、ファイル内容から精度を読み取ります。
    """
    data =[]

    iter_pattern =re .compile (r'iter(\d+)\.txt')

    acc_pattern =re .compile (r'After inference_outputs Correct: \d+/\d+ \((\d+\.\d+)%\)')

    if not os .path .isdir (directory_path ):
        print (f"Warning: Directory not found at {directory_path }")
        return []

    for filename in os .listdir (directory_path ):
        iter_match =iter_pattern .match (filename )
        if iter_match :
            iteration =int (iter_match .group (1 ))
            filepath =os .path .join (directory_path ,filename )
            try :
                with open (filepath ,'r')as f :
                    for line in f :
                        acc_match =acc_pattern .search (line )
                        if acc_match :
                            accuracy =float (acc_match .group (1 ))
                            data .append ((iteration ,accuracy ))
                            break 
            except IOError as e :
                print (f"Could not read file {filepath }: {e }")


    return sorted (data )

def generate_plot (data_series ,labels ,output_path ):
    """
    複数のデータ系列からグラフを生成し、画像として保存します。
    """
    plt .style .use ('seaborn-v0_8-whitegrid')
    fig ,ax =plt .subplots (figsize =(14 ,8 ))


    for i ,data in enumerate (data_series ):
        if not data :
            continue 
        iterations ,accuracies =zip (*data )
        ax .plot (iterations ,accuracies ,marker ='o',linestyle ='-',label =labels [i ])


    ax .set_title ('Inference Accuracy Trend',fontsize =18 )
    ax .set_xlabel ('Iteration',fontsize =12 )
    ax .set_ylabel ('Accuracy (%)',fontsize =12 )


    ax .yaxis .set_major_formatter (mticker .PercentFormatter ())


    ax .legend (title ='Dataset',loc ='lower right',fontsize ='medium')

    plt .tight_layout ()


    try :
        fig .savefig (output_path ,dpi =300 )
        print (f"Graph successfully saved to: {output_path }")
    except IOError as e :
        print (f"Failed to save graph: {e }")

def main ():
    """
    メイン関数
    """

    base_path ="/path/to/home/lltm/02_codeexec_etcot/scripts/ET-CoT/inference_outputs/Meta-Llama-3-8B-Instruct/3epoch"

    paths_to_plot =[
    f"{base_path }/LLTM-cruxeval-numeric-depth-val/LLTM-all-numeric-depth",
    f"{base_path }/LLTM-livecodebench-all-numeric-depth-val/LLTM-all-numeric-depth",
    f"{base_path }/LLTM-livecodebench-numeric-depth-val/LLTM-all-numeric-depth"
    ]


    labels =[path .split ('/')[-2 ]for path in paths_to_plot ]

    all_data =[]
    for path in paths_to_plot :
        print (f"Processing data from: {path }")
        all_data .append (parse_accuracy_data (path ))


    output_filename ="inference_trend_legend_bottom_right.png"


    generate_plot (all_data ,labels ,output_filename )

if __name__ =='__main__':
    main ()
