import numpy as np
import matplotlib.pyplot as plt

class MathematicalFunctions:
    """
    A class containing 24 mathematical functions with various characteristics
    including continuous, discontinuous, descending, and constant segment behaviors.
    Domain: [-100, 100], Codomain: discrete values from [-1.0, 0.05, 0.1, 0.25, 0.5, 1.0, 1.5, 2.5, 3.5, 5.0, 7.5, 10.0, 20.0, 30.0]
    """

    def __init__(self):
        """Initialize the class with all 24 mathematical functions"""
        # Define the allowed discrete codomain values
        self.allowed_values = np.array([-1.0, 0.05, 0.1, 0.25, 0.5, 1.0, 1.5, 2.5, 3.5, 5.0, 7.5, 10.0, 20.0, 30.0])
        self.functions = {
            'linear_ascending': self._linear_ascending,
            'linear_descending': self._linear_descending,
            'parabolic_upward': self._parabolic_upward,
            'parabolic_downward': self._parabolic_downward,
            'cubic_ascending': self._cubic_ascending,
            'cubic_descending': self._cubic_descending,
            'exponential_growth': self._exponential_growth,
            'exponential_decay': self._exponential_decay,
            'logarithmic_growth': self._logarithmic_growth,
            'logarithmic_decay': self._logarithmic_decay,
            'sine_wave': self._sine_wave,
            'cosine_wave': self._cosine_wave,
            'step_function_ascending': self._step_function_ascending,
            'step_function_descending': self._step_function_descending,
            'piecewise_mixed': self._piecewise_mixed,
            'sigmoidal_curve': self._sigmoidal_curve,
            'diurnal_cycle': self._diurnal_cycle,
            'weather_front': self._weather_front,
            'seasonal_trend': self._seasonal_trend,
            'humidity_buildup': self._humidity_buildup,
            'pressure_gradient': self._pressure_gradient,
            'humidity_buildup_descending': self._humidity_buildup_descending,
            'pressure_gradient_descending': self._pressure_gradient_descending,
            'specular_seasonal_trending': self._specular_seasonal_trending
        }

        self.function_names = list(self.functions.keys())

    def _map_to_allowed_values(self, y_values):
        """
        Map continuous y values to the nearest allowed discrete values.
        
        Parameters:
        -----------
        y_values : array-like
            Continuous y values to be mapped
            
        Returns:
        --------
        numpy.ndarray : Array of values mapped to allowed discrete values
        """
        y_values = np.asarray(y_values)
        mapped_values = np.zeros_like(y_values)
        
        for i, y in enumerate(y_values):
            # Find the closest allowed value
            distances = np.abs(self.allowed_values - y)
            closest_index = np.argmin(distances)
            mapped_values[i] = self.allowed_values[closest_index]
            
        return mapped_values

    def _linear_ascending(self, x):
        """Linear function - straight line ascending from allowed minimum to maximum"""
        continuous_y = 30 * (x + 100) / 200
        return self._map_to_allowed_values(continuous_y)

    def _linear_descending(self, x):
        """Linear function - straight line descending from allowed maximum to minimum"""
        continuous_y = 30 - 30 * (x + 100) / 200
        return self._map_to_allowed_values(continuous_y)

    def _parabolic_upward(self, x):
        """Quadratic function - parabola opening upward mapped to allowed values"""
        continuous_y = 30 * (x / 100) ** 2
        return self._map_to_allowed_values(continuous_y)

    def _parabolic_downward(self, x):
        """Quadratic function - inverted parabola opening downward mapped to allowed values"""
        continuous_y = 30 * (1 - (x / 100) ** 2)
        return self._map_to_allowed_values(continuous_y)

    def _cubic_ascending(self, x):
        """Cubic function - S-shaped curve ascending mapped to allowed values"""
        continuous_y = 15 + 15 * (x / 100) ** 3
        return self._map_to_allowed_values(continuous_y)

    def _cubic_descending(self, x):
        """Cubic function - S-shaped curve descending mapped to allowed values"""
        continuous_y = 15 - 15 * (x / 100) ** 3
        return self._map_to_allowed_values(continuous_y)

    def _exponential_growth(self, x):
        """Exponential function - exponential growth mapped to allowed values"""
        continuous_y = 30 * (1 - np.exp(-(x + 100) / 50))
        return self._map_to_allowed_values(continuous_y)

    def _exponential_decay(self, x):
        """Exponential function - exponential decay mapped to allowed values"""
        continuous_y = 30 * np.exp(-(x + 100) / 50)
        return self._map_to_allowed_values(continuous_y)

    def _logarithmic_growth(self, x):
        """Logarithmic function - logarithmic growth mapped to allowed values"""
        continuous_y = 30 * np.log(x + 101) / np.log(201)
        return self._map_to_allowed_values(continuous_y)

    def _logarithmic_decay(self, x):
        """Logarithmic function - logarithmic decay mapped to allowed values"""
        continuous_y = 30 * (1 - np.log(x + 101) / np.log(201))
        return self._map_to_allowed_values(continuous_y)

    def _sine_wave(self, x):
        """Trigonometric function - sine wave oscillating between allowed values"""
        continuous_y = 15 + 15 * np.sin(4 * np.pi * x / 200)
        return self._map_to_allowed_values(continuous_y)

    def _cosine_wave(self, x):
        """Trigonometric function - cosine wave oscillating between allowed values"""
        continuous_y = 15 + 15 * np.cos(4 * np.pi * x / 200)
        return self._map_to_allowed_values(continuous_y)

    def _step_function_ascending(self, x):
        """Step function - ascending discrete levels using only allowed values"""
        # Directly use allowed values for clean steps
        return np.where(x < -60, 0.05,
                        np.where(x < -20, 7.5,
                                 np.where(x < 20, 10.0,
                                          np.where(x < 60, 20.0, 30.0))))

    def _step_function_descending(self, x):
        """Step function - descending discrete levels using only allowed values"""
        # Directly use allowed values for clean steps
        return np.where(x < -60, 30.0,
                        np.where(x < -20, 20.0,
                                 np.where(x < 20, 10.0,
                                          np.where(x < 60, 7.5, 0.05))))

    def _piecewise_mixed(self, x):
        """Piecewise function - mix of ascending, constant, and descending segments using allowed values"""
        continuous_y = np.where(x < -60, 0.15 * (x + 100),  # ascending from 0
                                np.where(x < -20, 6,  # constant
                                         np.where(x < 20, 18,  # constant higher
                                                  np.where(x < 60, 30 - 0.3 * (x - 20), 18))))  # descending
        return self._map_to_allowed_values(continuous_y)

    def _sigmoidal_curve(self, x):
        """Sigmoidal curve - S-shaped transition mapped to allowed values"""
        continuous_y = 30 / (1 + np.exp(-x / 25))
        return self._map_to_allowed_values(continuous_y)

    def _diurnal_cycle(self, x):
        """Diurnal cycle - realistic daily water vapor variation mapped to allowed values"""
        daily_cycle = np.sin(2 * np.pi * (x + 100) / 100 + np.pi / 4)
        continuous_y = 12 + 8 * daily_cycle + 3 * np.sin(4 * np.pi * (x + 100) / 100)
        return self._map_to_allowed_values(continuous_y)

    def _weather_front(self, x):
        """Weather front passage - gradual increase then gradual decrease mapped to allowed values"""
        front_center = 0
        width = 60
        distance_from_center = np.abs(x - front_center)
        front_effect = 25 * np.exp(-(distance_from_center / width) ** 2)
        continuous_y = 5 + front_effect
        return self._map_to_allowed_values(continuous_y)

    def _seasonal_trend(self, x):
        """Seasonal trend - long-term gradual changes mapped to allowed values"""
        seasonal_base = 15 + 10 * np.sin(np.pi * (x + 100) / 200)
        variability = 3 * np.sin(3 * np.pi * (x + 100) / 200)
        continuous_y = seasonal_base + variability
        return self._map_to_allowed_values(continuous_y)

    def _humidity_buildup(self, x):
        """Humidity buildup - gradual accumulation with occasional releases mapped to allowed values"""
        buildup = 8 + 0.15 * (x + 100)
        release_cycles = 5 * np.sin(8 * np.pi * (x + 100) / 200) * np.exp(-((x + 50) / 80) ** 2)
        continuous_y = buildup - release_cycles
        continuous_y = np.clip(continuous_y, 0, 30)
        return self._map_to_allowed_values(continuous_y)

    def _pressure_gradient(self, x):
        """Pressure gradient effect - smooth transitions between pressure zones mapped to allowed values"""
        zone1 = 8 + 3 * np.tanh((x + 60) / 20)
        zone2 = 18 + 5 * np.tanh((x - 20) / 25)
        weight = 0.5 + 0.5 * np.tanh(x / 40)
        continuous_y = zone1 * (1 - weight) + zone2 * weight
        return self._map_to_allowed_values(continuous_y)

    def _humidity_buildup_descending(self, x):
        """Humidity buildup descending - gradual decrease with occasional buildups mapped to allowed values"""
        # Start high and decrease
        buildup = 30 - 8 - 0.15 * (x + 100)
        # Add occasional buildup cycles (opposite of releases)
        buildup_cycles = 5 * np.sin(8 * np.pi * (x + 100) / 200) * np.exp(-((x + 50) / 80) ** 2)
        continuous_y = buildup + buildup_cycles
        continuous_y = np.clip(continuous_y, 0, 30)
        return self._map_to_allowed_values(continuous_y)

    def _pressure_gradient_descending(self, x):
        """Pressure gradient descending - smooth transitions from high to low pressure zones mapped to allowed values"""
        # Inverted pressure gradient - start high, end low
        zone1 = 22 - 3 * np.tanh((x + 60) / 20)  # High pressure zone
        zone2 = 12 - 5 * np.tanh((x - 20) / 25)  # Lower pressure zone
        weight = 0.5 + 0.5 * np.tanh(x / 40)
        continuous_y = zone1 * (1 - weight) + zone2 * weight
        return self._map_to_allowed_values(continuous_y)

    def _specular_seasonal_trending(self, x):
        """Specular seasonal trending - mirror-like seasonal patterns with reflective symmetry mapped to allowed values"""
        # Create a complex pattern with multiple harmonics and reflective properties
        base_trend = 15 + 8 * np.cos(2 * np.pi * (x + 100) / 200)  # Main seasonal cycle
        
        # Add specular (mirror-like) components
        specular_component = 5 * np.sin(np.pi * (x + 100) / 100) * np.cos(3 * np.pi * (x + 100) / 200)
        
        # Add reflection pattern around center
        reflection_pattern = 3 * np.exp(-((x - 0) / 40) ** 2) * np.sin(6 * np.pi * x / 200)
        
        # Combine all components
        continuous_y = base_trend + specular_component + reflection_pattern
        continuous_y = np.clip(continuous_y, 0, 30)
        
        return self._map_to_allowed_values(continuous_y)

    def get(self, function_name, x_start, x_end, length):
        """
        Get y values from a specific function over a range of x values.

        Parameters:
        -----------
        function_name : str
            Name of the function to use (from self.function_names)
        x_start : float
            Starting x value
        x_end:
            Ending x value
        length : int
            Number of y values to return (determines x range)

        Returns:
        --------
        dict : Dictionary containing:
            - 'x_values': array of x values from x_start to x_start + length - 1
            - 'y_values': corresponding y values from the chosen function
            - 'function_used': name of the function used

        Raises:
        -------
        ValueError: If function_name is not valid or x values exceed domain bounds
        """

        # Validate function name
        if function_name not in self.functions:
            available_functions = ", ".join(self.function_names)
            raise ValueError(f"Function '{function_name}' not found. Available functions: {available_functions}")

        # Create x values array
        x_values = np.linspace(x_start, x_end, length, dtype=float)

        # Check domain bounds
        if np.any(x_values < -100) or np.any(x_values > 100):
            raise ValueError(f"X values must be within domain [-100, 100]. "
                             f"Requested range: [{x_start}, {x_start + length - 1}]")

        # Get the function and calculate y values
        selected_function = self.functions[function_name]
        y_values = selected_function(x_values)

        return {
            'x_values': x_values,
            'y_values': y_values,
            'function_used': function_name
        }

    def list_functions(self):
        """Return a list of all available function names"""
        return self.function_names.copy()

