from tensorboard.backend.event_processing import event_accumulator

import os
from shutil import copy2
from re import search as RSearch
import pandas as pd
from ast import literal_eval as LEval

weights_dir = 'logs/weights/'

def find_biggest_tensorboard(tensordir):
    try:
        files = [f for f in os.listdir(tensordir) if f.endswith('.0')]
        if not files:
            print("No files with the '.0' extension found!")
            return

        max_size = 0
        biggest_file = ""

        for file in files:
            file_path = os.path.join(tensordir, file)
            if os.path.isfile(file_path):
                file_size = os.path.getsize(file_path)
                if file_size > max_size:
                    max_size = file_size
                    biggest_file = file

        return biggest_file

    except FileNotFoundError:
        print("Couldn't find your model!")
        return

def main(model_name, save_freq, lastmdls):
    global lowestval_weight_dir, scl

    tensordir = os.path.join('logs', model_name)
    lowestval_weight_dir = os.path.join(tensordir, "lowestvals")
    
    latest_file = find_biggest_tensorboard(tensordir)
    
    if latest_file is None:
        print("Couldn't find a valid tensorboard file!")
        return
    
    tfile = os.path.join(tensordir, latest_file)
    
    ea = event_accumulator.EventAccumulator(tfile,
        size_guidance={
        event_accumulator.COMPRESSED_HISTOGRAMS: 500,
        event_accumulator.IMAGES: 4,
        event_accumulator.AUDIO: 4,
        event_accumulator.SCALARS: 0,
        event_accumulator.HISTOGRAMS: 1,
    })

    ea.Reload()
    ea.Tags()

    scl = ea.Scalars('loss/g/total')

    listwstep = {}
    
    for val in scl:
        if (val.step // save_freq) * save_freq in [val.step for val in scl]:
            listwstep[float(val.value)] = (val.step // save_freq) * save_freq

    lowest_vals = sorted(listwstep.keys())[:lastmdls]

    sorted_dict = {value: step for value, step in listwstep.items() if value in lowest_vals}
    
    return sorted_dict

def selectweights(model_name, file_dict, weights_dir, lowestval_weight_dir):
    os.makedirs(lowestval_weight_dir, exist_ok=True)
    logdir = []
    files = []
    lbldict = {
        'Values': {},
        'Names': {}
    }
    weights_dir_path = os.path.join(weights_dir, "")
    low_val_path = os.path.join(os.getcwd(), os.path.join(lowestval_weight_dir, ""))
    
    try:
        file_dict = LEval(file_dict)
    except Exception as e: 
        print(f"Error! {e}")
        return f"Couldn't load tensorboard file! {e}"
    
    weights = [f for f in os.scandir(weights_dir)]
    for key, value in file_dict.items():
        pattern = fr"^{model_name}_.*_s{value}\.pth$"
        matching_weights = [f.name for f in weights if f.is_file() and RSearch(pattern, f.name)]
        for weight in matching_weights:
            source_path = weights_dir_path + weight
            destination_path = os.path.join(lowestval_weight_dir, weight)
            
            copy2(source_path, destination_path)

            logdir.append(f"File = {weight} Value: {key}, Step: {value}")

            lbldict['Names'][weight] = weight
            lbldict['Values'][weight] = key

            files.append(low_val_path + weight)

            print(f"File = {weight} Value: {key}, Step: {value}")

            yield ('\n'.join(logdir), files, pd.DataFrame(lbldict))
            

    return ''.join(logdir), files, pd.DataFrame(lbldict)
    

if __name__ == "__main__":
    model = str(input("Enter the name of the model: "))
    sav_freq = int(input("Enter save frequency of the model: "))
    ds = main(model, sav_freq)
    
    if ds: selectweights(model, ds, weights_dir, lowestval_weight_dir)