import pickle
import numpy as np
import matplotlib.pyplot as plt
import glob
import os
from datetime import datetime

def plot_results():
    filename = r"data\Anchored_AMDP\kl\.pkl"

    with open(filename, 'rb') as f:
        sample_error_data = pickle.load(f)
    
    plt.figure(figsize=(10, 6))

    n_samples_list = sorted(sample_error_data.keys())
    errors = [sample_error_data[n_sample] for n_sample in n_samples_list]
    plt.plot(n_samples_list, errors, marker='o', linestyle='none', color='b', label='Error vs Samples')

    log_n_samples = np.log(n_samples_list)
    log_errors = np.log(errors)
    coefficients = np.polyfit(log_n_samples, log_errors, 1)
    slope, intercept = coefficients
    log_n_samples_fit = np.linspace(min(log_n_samples), max(log_n_samples), 100)
    log_errors_fit = slope * log_n_samples_fit + intercept

    n_sample_fit = np.exp(log_n_samples_fit)
    errors_fit = np.exp(log_errors_fit)
    plt.plot(n_sample_fit, errors_fit, linestyle='--', color='red', label=f'Fitted Line (Slope = {slope:.2f})')

    x_min, x_max = min(n_samples_list), max(n_samples_list)
    y_min, y_max = min(errors), max(errors)

    C = np.sqrt((errors[0] * errors[-1])) * (n_samples_list[0] ** 0.5)
    x_line = np.linspace(x_min, x_max, 100)
    y_line = C * (x_line ** (-0.5))

    # plot the line with slope -1/2
    plt.plot(x_line, y_line, linestyle='--', color='black', label='Slope = -1/2')

    plt.xscale('log')
    plt.yscale('log')

    plt.title('Error vs. Sample (Log Scale)')
    plt.xlabel('Sample (Log Scale)')
    plt.ylabel('Error (Log Scale)')

    plt.grid(True, which="both", ls="--")
    plt.legend(loc='upper right')
    # plt.show()
    plt.savefig(r'data\Anchored_AMDP\kl\.png')
    print(sample_error_data.keys())
if __name__ == '__main__':
    plot_results()