import os
import csv
from fine_tuning.wrappers import NeuroRVQWrapper

class CSVLogger():
    def __init__(self, output_dir, ex_id):
        self.log_dir = os.path.join(output_dir, f"{ex_id}_log")
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)
        self._files = set()
    
    def report_scalar(self, title, series, value, iteration):
        '''
        Mimics clearml report_scalar() function to log values to CSV file
        '''
        if 'train' in series:
            filepath = os.path.join(self.log_dir, f"{title}_train.csv")
        else:
            filepath = os.path.join(self.log_dir, f"{title}_val.csv")

        write_header = filepath not in self._files

        with open(filepath, mode="a", newline="") as f:
            writer = csv.writer(f)
            if 'MEAN' in title:
                if write_header:
                    writer.writerow(["Series", "Iteration", "Value"])
                    self._files.add(filepath)
                writer.writerow([series, iteration, value])
            else:
                if write_header:
                    writer.writerow(["Fold", "Iteration", "Value"])
                    self._files.add(filepath)
                writer.writerow([series.split(' ')[-1], iteration, value])

def get_logger():
    logger = CSVLogger("results", 0)
    return logger

def get_model(ch_names, n_times, n_outputs, args, foundation_model, train_head_only=False):
    """
    Returns: FinetuningWrapper for the specified model
    """
    return NeuroRVQWrapper(
        n_time=n_times,
        ch_names=ch_names,
        n_outputs=n_outputs,
        train_head_only=train_head_only,
        args = args,
        foundation_model = foundation_model
    )