import numpy as np
import pandas as pd
from scipy.optimize import curve_fit

def cosinor_model(t, MESOR, Amplitude, Acrophase):
    """Cosine function for circadian rhythm fitting."""
    return MESOR + Amplitude * np.cos(2 * np.pi * t / 24 + Acrophase)

def cosinor_analysis(time_stamps, signal):
    """
    Perform cosinor analysis to extract MESOR, Amplitude, and Acrophase.

    Args:
        time_stamps (array-like): Time values (should be datetime-like).
        signal (array-like): Temperature values.

    Returns:
        MESOR (float): Mean baseline level.
        Amplitude (float): Strength of oscillation.
        Acrophase (float): Time of peak (in hours).
    """
    # 🔹 Convert time_stamps to Pandas Series and remove NaN values
    time_stamps = pd.Series(time_stamps)
    signal = pd.Series(signal)

    # 🔹 Ensure time_stamps are datetime
    time_stamps = pd.to_datetime(time_stamps, errors='coerce')

    # 🔹 Drop NaNs from both arrays
    valid_idx = (~time_stamps.isna()) & (~signal.isna())
    time_stamps = time_stamps[valid_idx]
    signal = signal[valid_idx]

    # 🔹 Convert datetime to hours since first measurement
    time_hours = (time_stamps - time_stamps.min()).dt.total_seconds() / 3600.0

    # 🔹 Ensure signal is a NumPy array
    signal = np.array(signal, dtype=np.float64)

    # 🔹 Check if the lengths match
    if len(time_hours) != len(signal):
        raise ValueError(f"Mismatch: time_hours ({len(time_hours)}) != signal ({len(signal)})")

    # Initial estimates (MESOR = mean, Amplitude = half range, Acrophase = 0)
    initial_guess = [np.mean(signal), (np.max(signal) - np.min(signal)) / 2, 0]
    #import IPython; IPython.embed()
    try:
        # Fit the cosine function
        popt, _ = curve_fit(cosinor_model, time_hours, signal, p0=initial_guess)
        MESOR, Amplitude, Acrophase = popt

        # Convert Acrophase from radians to hours
        Acrophase_hours = (Acrophase * 24) / (2 * np.pi)

        return MESOR, Amplitude, Acrophase_hours

    except RuntimeError:
        # Fitting failed, return NaN
        return np.nan, np.nan, np.nan
