"""
utils.py stores some general purpose function and classes.
"""

import os 
import numpy as np 
import logging 
import time 
from joblib import cpu_count 
from shutil import rmtree, copy, copytree


##


class TimerError(Exception):
    """
    A custom exception used to report errors in use of Timer class.
    """

class Timer:
    """
    A custom Timer class.
    """
    def __init__(self):
        self._start_time = None

    def start(self):
        """
        Start a new timer.
        """
        if self._start_time is not None:
            raise TimerError(f"Timer is running. Use .stop() to stop it")
        self._start_time = time.perf_counter()

    def stop(self):
        """
        Stop the timer, and report the elapsed time.
        """
        if self._start_time is None:
            raise TimerError(f"Timer is not running. Use .start() to start it")
        elapsed_time = time.perf_counter() - self._start_time

        if elapsed_time > 100:
            unit = 'min'
            elapsed_time = elapsed_time / 60
        elif elapsed_time > 1000:
            unit = 'h'
            elapsed_time = elapsed_time / 3600
        else:
            unit = 's'

        self._start_time = None

        return f'{round(elapsed_time, 2)} {unit}'


##


def make_folder(path, name, overwrite=True):
    """
    A function to create a new {name} folder at the {path} path.
    """
    os.chdir(path)
    if not os.path.exists(name) or overwrite:
        rmtree(os.path.join(path, name), ignore_errors=True)
        os.mkdir(name)
    else:
        pass


##


def set_logger(path_runs, name, mode='w'):
    """
    A function to open a logs.txt file for a certain script, writing its trace at path_main/runs/step/.
    """
    logger = logging.getLogger("Cellula_logs")
    handler = logging.FileHandler(os.path.join(path_runs, name), mode=mode)
    handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
    logger.setLevel(logging.INFO)
    logger.addHandler(handler)

    return logger


##


def chunker(n):
    """
    Create an np.array of starting indeces for parallel computation.
    """
    n_jobs = cpu_count()
    starting_indeces = np.zeros(n_jobs + 1, dtype=int)
    quotient = n // n_jobs
    remainder = n % n_jobs

    for i in range(n_jobs):
        starting_indeces[i+1] = starting_indeces[i] + quotient + (1 if i < remainder else 0)

    return starting_indeces


##


def run_command(func, *args, verbose=False, **kwargs):
    """
    Helper function caller.
    """
    if verbose:
        print(f'{func.__name__} called with *args {args} and **kwargs {kwargs}')

    t = Timer()
    t.start()
    out = func(*args, **kwargs)
    if verbose:
        print(f'Elapsed time: {t.stop()}')
    
    return out

##


def update_params(d_original, d_passed):
    for k in d_passed:
        if k in d_original:
            pass
        else:
            print(f'{k}:{d_passed[k]} kwargs added...')
        d_original[k] = d_passed[k]
        
    return d_original


##


def save_best_pdf_quality(fig, figsize, path, name, dpi=1000):

    import matplotlib.pyplot as plt
    import io
    from PIL import Image
    import os

    # Path to save the final PDF
    path_pdf = os.path.join(path, name)

    # Step 1: Render your figure as a high-DPI PNG in memory
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=dpi)
    buf.seek(0)

    # Step 2: Create a new PDF with the PNG embedded
    fig2, ax = plt.subplots(figsize=figsize)  # Match your original figure size
    img = Image.open(buf)
    ax.imshow(img)
    ax.axis('off')

    # Step 3: Save to PDF
    fig2.savefig(path_pdf, bbox_inches='tight', dpi=dpi)
    buf.close()
    plt.close(fig2)


##