import time

import numpy as np
from datetime import datetime, timedelta


def _simulate_temperature_data(
    start_temp=25,  # Room temperature (ºC)
    first_target=-70,  # First cooling target (ºC)
    second_target=-40,  # Heating target (ºC)
    final_target=-180,  # Final cooling target (ºC)
    total_time=720,  # Total time in minutes (12 hours)
    sample_rate=10,  # Seconds between measurements
) -> tuple[list[datetime], list[float]]:
    """
    Simulate temperature sensor data with a plateau phase and a cooling period.

    The default arguments will generate temperature data for twelve hours cooling
    down from room temperature to -180ºC. We introduced two plateau of one hour
    duration, (1) the first plateau at -70ºC starts after one hour of cooling and
    (2) the second plateau at -40ºC starts after heating for one hour and twenty
    minutes. After the second plateau we cool down until -180ºC.

    This code was generated by Claude 3.5 Sonnet with small adjustments.

    Returns:
        - timestamp and temperature as a tuple.
    """
    # Calculate number of samples
    n_samples = int((total_time * 60) / sample_rate)

    # Create time array in minutes
    time = np.linspace(0, total_time, n_samples)

    # Initialize temperature array
    temp = np.zeros(n_samples)

    # Define phase transition points (in minutes)
    phase1_end = 60  # First cooling phase
    phase2_end = 120  # First plateau
    phase3_end = 200  # Heating phase
    phase4_end = 260  # Second plateau

    # Find indices for different phases
    idx1 = np.where(time >= phase1_end)[0][0]
    idx2 = np.where(time >= phase2_end)[0][0]
    idx3 = np.where(time >= phase3_end)[0][0]
    idx4 = np.where(time >= phase4_end)[0][0]

    # Phase 1: Initial cooling to -70°C (exponential)
    initial_time = time[:idx1]
    temp[:idx1] = start_temp + (first_target - start_temp) * (
        1 - np.exp(-initial_time / 15)  # Time constant for first cooling
    )

    # Phase 2: First plateau at -70°C
    temp[idx1:idx2] = first_target

    # Phase 3: Heating to -40°C (exponential approach)
    heating_time = time[idx2:idx3] - time[idx2]
    temp[idx2:idx3] = first_target + (second_target - first_target) * (
        1 - np.exp(-heating_time / 10)  # Time constant for heating
    )

    # Phase 4: Second plateau at -40°C
    temp[idx3:idx4] = second_target

    # Phase 5: Final cooling to -180°C
    final_time = time[idx4:] - time[idx4]
    temp[idx4:] = second_target + (final_target - second_target) * (
        1 - np.exp(-final_time / 100)  # Time constant for final cooling
    )

    # Add different levels of noise for each phase
    base_noise = np.random.normal(0, 0.2, n_samples)

    # Add occasional spikes
    spike_locations = np.random.choice(
        n_samples, size=int(n_samples * 0.005), replace=False
    )
    spike_noise = np.zeros(n_samples)
    spike_noise[spike_locations] = np.random.normal(0, 3, size=len(spike_locations))

    # Add slow drift
    drift = np.sin(np.linspace(0, 8 * np.pi, n_samples)) * 0.8

    # Combine all noise components
    total_noise = base_noise + spike_noise + drift
    temp += total_noise

    # Create timestamps
    start_time = datetime.now()
    timestamps = [
        start_time + timedelta(seconds=i * sample_rate) for i in range(n_samples)
    ]

    return timestamps, temp.tolist()


class SimulatedTemperature:
    """
    This is an iterator over simulated temperature data.

    Generate simulate temperature sensor data with a plateau phase and a cooling period.

    The default arguments will generate temperature data for twelve hours cooling
    down from room temperature to -180ºC. We introduced two plateau of one hour
    duration, (1) the first plateau at -70ºC starts after one hour of cooling and
    (2) the second plateau at -40ºC starts after heating for one hour and twenty
    minutes. After the second plateau we cool down until -180ºC.

    """

    def __init__(self):
        self.timestamps: list[datetime]
        self.temperatures: list[float]

        self.timestamps, self.temperatures = _simulate_temperature_data()

        self.count: int = len(self.timestamps)
        """The number of data points."""
        self.idx: int = 0
        """A pointer in the list of data points, the next data point to return."""

    def __iter__(self):
        return self

    def __next__(self):
        x, y = self.timestamps[self.idx], self.temperatures[self.idx]
        if self.idx >= self.count:
            raise StopIteration
        self.idx += 1
        return x, y


if __name__ == "__main__":
    # for timestamp, temperature in SimulatedTemperature():
    #     print(f"{timestamp = }, {temperature = }")
    #     time.sleep(1.0)

    import matplotlib.pyplot as plt

    x, y = _simulate_temperature_data()

    plt.plot(x, y)
    plt.show()
