import numpy as np
import matplotlib.pyplot as plt
from typing import Callable, List, Tuple, Optional, Dict, Any, Union
import dataclasses
from enum import Enum
import inspect

# ====================== ТИПЫ ЗАВИСИМОСТЕЙ ======================

class DependencyType(Enum):
    """Типы зависимостей между параметрами"""
    INDEPENDENT = "independent"  # Независимый параметр
    DEPENDENT = "dependent"      # Зависимый параметр
    EXPRESSION = "expression"    # Выражение, использующее другие параметры

# ====================== КЛАССЫ ДЛЯ ПАРАМЕТРОВ ======================

@dataclasses.dataclass
class Parameter:
    """Базовый класс для параметра"""
    name: str
    param_type: DependencyType
    description: str = ""
    
    def evaluate(self, context: Dict[str, Any]) -> np.ndarray:
        """Вычисляет значение параметра в заданном контексте"""
        raise NotImplementedError

@dataclasses.dataclass
class IndependentParameter(Parameter):
    """Независимый параметр (базовый)"""
    formula: Callable
    x_range: Tuple[float, float]
    num_points: int = 1000
    color: str = 'blue'
    line_style: str = '-'
    
    def __post_init__(self):
        self.param_type = DependencyType.INDEPENDENT
    
    def evaluate(self, context: Dict[str, Any]) -> np.ndarray:
        """Вычисляет значение независимого параметра"""
        x = context.get('x', np.linspace(self.x_range[0], self.x_range[1], self.num_points))
        return self.formula(x)

@dataclasses.dataclass
class DependentParameter(Parameter):
    """Зависимый параметр (выражается через другие параметры)"""
    expression: Callable  # Функция, которая использует другие параметры
    dependencies: List[str]  # Имена параметров, от которых зависит
    color: str = 'green'
    line_style: str = '--'
    
    def __post_init__(self):
        self.param_type = DependencyType.DEPENDENT
    
    def evaluate(self, context: Dict[str, Any]) -> np.ndarray:
        """Вычисляет значение зависимого параметра"""
        # Собираем значения зависимостей из контекста
        dep_values = {}
        for dep_name in self.dependencies:
            if dep_name not in context:
                raise ValueError(f"Зависимость '{dep_name}' не найдена в контексте")
            dep_values[dep_name] = context[dep_name]
        
        # Вычисляем выражение
        return self.expression(**dep_values)

# ====================== КЛАСС ДЛЯ ГАРМОНИК ======================

class HarmonicOscillation:
    """Класс для описания гармонического колебания"""
    def __init__(self, amplitude: Union[float, str], 
                 frequency: Union[float, str], 
                 phase: Union[float, str] = 0,
                 amplitude_depends_on: Optional[List[str]] = None,
                 frequency_depends_on: Optional[List[str]] = None,
                 phase_depends_on: Optional[List[str]] = None):
        """
        Параметры гармоники могут быть как числами, так и именами параметров модели
        """
        self.amplitude = amplitude
        self.frequency = frequency
        self.phase = phase
        
        # Отслеживаем зависимости
        self.amplitude_depends_on = amplitude_depends_on or []
        self.frequency_depends_on = frequency_depends_on or []
        self.phase_depends_on = phase_depends_on or []
    
    def get_dependencies(self) -> List[str]:
        """Возвращает все зависимости гармоники"""
        return (self.amplitude_depends_on + 
                self.frequency_depends_on + 
                self.phase_depends_on)
    
    def evaluate(self, t: np.ndarray, context: Dict[str, Any]) -> np.ndarray:
        """Вычисляет значение гармонического колебания с учетом зависимостей"""
        
        # Получаем значения амплитуды, частоты и фазы
        amp = self._get_value(self.amplitude, context)
        freq = self._get_value(self.frequency, context)
        phase = self._get_value(self.phase, context)
        
        return amp * np.sin(2 * np.pi * freq * t + phase)
    
    def _get_value(self, param: Any, context: Dict[str, Any]) -> float:
        """Получает значение параметра (число или из контекста)"""
        if isinstance(param, (int, float)):
            return param
        elif isinstance(param, str) and param in context:
            # Если параметр - строка, берем соответствующее значение из контекста
            val = context[param]
            # Если это массив, берем среднее или первое значение
            if isinstance(val, np.ndarray):
                return np.mean(val)  # или val[0] в зависимости от логики
            return float(val)
        else:
            raise ValueError(f"Не удалось получить значение для {param}")
    
    def __repr__(self):
        return f"Harmonic(A={self.amplitude}, f={self.frequency}, φ={self.phase})"

# ====================== ОСНОВНОЙ КЛАСС МОДЕЛИ ======================

class MathematicalModel:
    """
    Модель с поддержкой зависимых параметров
    """
    
    def __init__(self):
        self.parameters: Dict[str, Parameter] = {}
        self.harmonics: List[HarmonicOscillation] = []
        self.main_formula: Optional[Callable] = None
        self.evaluation_context: Dict[str, Any] = {}
    
    def add_independent_parameter(self, name: str, formula: Callable, 
                                   x_range: Tuple[float, float], 
                                   num_points: int = 1000,
                                   color: str = 'blue',
                                   line_style: str = '-',
                                   description: str = "") -> None:
        """Добавляет независимый параметр"""
        self.parameters[name] = IndependentParameter(
            name=name,
            formula=formula,
            x_range=x_range,
            num_points=num_points,
            color=color,
            line_style=line_style,
            description=description
        )
    
    def add_dependent_parameter(self, name: str, expression: Callable,
                                  dependencies: List[str],
                                  color: str = 'green',
                                  line_style: str = '--',
                                  description: str = "") -> None:
        """Добавляет зависимый параметр"""
        # Проверяем, что все зависимости существуют
        for dep in dependencies:
            if dep not in self.parameters:
                raise ValueError(f"Зависимость '{dep}' не найдена среди параметров")
        
        self.parameters[name] = DependentParameter(
            name=name,
            expression=expression,
            dependencies=dependencies,
            color=color,
            line_style=line_style,
            description=description
        )
    
    def add_harmonic(self, amplitude: Union[float, str], 
                     frequency: Union[float, str], 
                     phase: Union[float, str] = 0,
                     amplitude_depends_on: Optional[List[str]] = None,
                     frequency_depends_on: Optional[List[str]] = None,
                     phase_depends_on: Optional[List[str]] = None) -> None:
        """Добавляет гармоническое колебание с возможными зависимостями"""
        
        harmonic = HarmonicOscillation(
            amplitude=amplitude,
            frequency=frequency,
            phase=phase,
            amplitude_depends_on=amplitude_depends_on,
            frequency_depends_on=frequency_depends_on,
            phase_depends_on=phase_depends_on
        )
        
        # Проверяем, что все зависимости существуют
        for dep in harmonic.get_dependencies():
            if dep not in self.parameters:
                raise ValueError(f"Зависимость '{dep}' для гармоники не найдена среди параметров")
        
        self.harmonics.append(harmonic)
    
    def set_main_formula(self, formula: Callable) -> None:
        """Устанавливает основную формулу"""
        self.main_formula = formula
    
    def get_parameter_dependencies(self, param_name: str) -> List[str]:
        """Возвращает список зависимостей для параметра"""
        param = self.parameters.get(param_name)
        if isinstance(param, DependentParameter):
            return param.dependencies
        return []
    
    def get_all_dependencies(self) -> Dict[str, List[str]]:
        """Возвращает словарь всех зависимостей"""
        dependencies = {}
        for name, param in self.parameters.items():
            if isinstance(param, DependentParameter):
                dependencies[name] = param.dependencies
        return dependencies
    
    def evaluate_parameters(self, x: np.ndarray) -> Dict[str, np.ndarray]:
        """
        Вычисляет все параметры с учетом зависимостей.
        Использует топологическую сортировку для правильного порядка вычисления.
        """
        results = {'x': x}
        
        # Функция для топологической сортировки
        def topological_sort():
            visited = set()
            order = []
            
            def dfs(param_name):
                if param_name in visited:
                    return
                visited.add(param_name)
                
                param = self.parameters[param_name]
                if isinstance(param, DependentParameter):
                    for dep in param.dependencies:
                        if dep in self.parameters:
                            dfs(dep)
                
                order.append(param_name)
            
            for name in self.parameters:
                if name not in visited:
                    dfs(name)
            
            return order
        
        # Вычисляем параметры в правильном порядке
        eval_order = topological_sort()
        
        for param_name in eval_order:
            param = self.parameters[param_name]
            results[param_name] = param.evaluate(results)
        
        return results
    
    def sum_harmonics(self, t: np.ndarray, context: Dict[str, Any]) -> np.ndarray:
        """Вычисляет сумму всех гармонических колебаний с учетом зависимостей"""
        if not self.harmonics:
            return np.zeros_like(t)
        
        result = np.zeros_like(t)
        for h in self.harmonics:
            result += h.evaluate(t, context)
        
        return result
    
    def evaluate_main(self, x: np.ndarray, **kwargs) -> np.ndarray:
        """
        Вычисляет основную формулу с заданным x и дополнительными параметрами
        """
        if self.main_formula is None:
            raise ValueError("Основная формула не установлена")
        
        # Вычисляем все параметры
        context = self.evaluate_parameters(x)
        context.update(kwargs)
        
        # Добавляем сумму гармоник в контекст
        context['harmonic_sum'] = lambda t: self.sum_harmonics(t, context)
        
        # Получаем аргументы функции
        sig = inspect.signature(self.main_formula)
        
        # Подготавливаем аргументы для вызова
        call_args = {}
        for param_name in sig.parameters:
            if param_name in context:
                call_args[param_name] = context[param_name]
            elif param_name == 'x':
                call_args['x'] = x
        
        return self.main_formula(**call_args)

# ====================== УЛУЧШЕННЫЙ ВИЗУАЛИЗАТОР ======================

class ModelVisualizer:
    """Класс для визуализации модели с зависимостями"""
    
    def __init__(self, model: MathematicalModel):
        self.model = model
        self.figures = {}
    
    def plot_parameter(self, param_name: str, 
                       x_range: Optional[Tuple[float, float]] = None,
                       figsize: Tuple[int, int] = (10, 6),
                       title: Optional[str] = None,
                       show_dependencies: bool = True) -> plt.Figure:
        """
        Строит график для параметра с учетом его зависимостей
        """
        if param_name not in self.model.parameters:
            raise ValueError(f"Параметр '{param_name}' не найден")
        
        param = self.model.parameters[param_name]
        
        # Определяем диапазон x
        if x_range is None:
            if isinstance(param, IndependentParameter):
                x_range = param.x_range
            else:
                # Для зависимых параметров нужно определить разумный диапазон
                x_range = (0, 10)  # По умолчанию
        
        x = np.linspace(x_range[0], x_range[1], 1000)
        
        # Вычисляем значение параметра
        context = self.model.evaluate_parameters(x)
        y = context[param_name]
        
        # Создаем фигуру
        fig, ax = plt.subplots(figsize=figsize)
        
        # Строим основной график
        ax.plot(x, y, color=param.color, linestyle=param.line_style, 
                linewidth=2, label=param_name)
        
        # Если нужно показать зависимости
        if show_dependencies and isinstance(param, DependentParameter):
            for dep_name in param.dependencies:
                if dep_name in context:
                    dep_y = context[dep_name]
                    # Нормализуем для отображения на том же графике
                    dep_y_normalized = (dep_y - np.min(dep_y)) / (np.max(dep_y) - np.min(dep_y) + 1e-10)
                    ax.plot(x, dep_y_normalized, '--', alpha=0.5, 
                           label=f"{dep_name} (норм.)")
        
        ax.grid(True, alpha=0.3)
        ax.set_title(title or f"График параметра: {param_name}\n{param.description}", fontsize=14)
        ax.set_xlabel("x", fontsize=12)
        ax.set_ylabel("Значение", fontsize=12)
        ax.legend(loc='best')
        
        self.figures[f"param_{param_name}"] = fig
        return fig
    
    def plot_parameter_with_dependencies(self, param_name: str,
                                         figsize: Tuple[int, int] = (14, 10)) -> plt.Figure:
        """
        Строит подробный график параметра и всех его зависимостей
        """
        if param_name not in self.model.parameters:
            raise ValueError(f"Параметр '{param_name}' не найден")
        
        param = self.model.parameters[param_name]
        
        # Собираем все зависимости (рекурсивно)
        def get_all_deps(name, deps_set):
            deps = self.model.get_parameter_dependencies(name)
            for dep in deps:
                if dep not in deps_set:
                    deps_set.add(dep)
                    get_all_deps(dep, deps_set)
            return deps_set
        
        all_deps = list(get_all_deps(param_name, set()))
        
        # Определяем диапазон x
        x_range = (0, 10)
        if isinstance(param, IndependentParameter):
            x_range = param.x_range
        elif all_deps:
            # Пытаемся найти независимый параметр среди зависимостей
            for dep in all_deps:
                p = self.model.parameters[dep]
                if isinstance(p, IndependentParameter):
                    x_range = p.x_range
                    break
        
        x = np.linspace(x_range[0], x_range[1], 1000)
        context = self.model.evaluate_parameters(x)
        
        # Создаем подграфики
        n_plots = len(all_deps) + 1
        cols = min(3, n_plots)
        rows = (n_plots + cols - 1) // cols
        
        fig, axes = plt.subplots(rows, cols, figsize=figsize)
        axes = axes.flatten() if n_plots > 1 else [axes]
        
        # График целевого параметра
        axes[0].plot(x, context[param_name], color=param.color, 
                    linestyle=param.line_style, linewidth=2.5)
        axes[0].grid(True, alpha=0.3)
        axes[0].set_title(f"{param_name} (целевой)", fontsize=12)
        axes[0].set_xlabel("x")
        
        # Графики зависимостей
        for i, dep_name in enumerate(all_deps, 1):
            if i < len(axes):
                dep_param = self.model.parameters[dep_name]
                axes[i].plot(x, context[dep_name], color=dep_param.color,
                            linestyle=dep_param.line_style, linewidth=2)
                axes[i].grid(True, alpha=0.3)
                axes[i].set_title(f"{dep_name}\n{getattr(dep_param, 'description', '')}", fontsize=10)
                axes[i].set_xlabel("x")
        
        # Скрываем лишние подграфики
        for j in range(len(all_deps) + 1, len(axes)):
            axes[j].set_visible(False)
        
        plt.tight_layout()
        self.figures[f"param_{param_name}_with_deps"] = fig
        return fig
    
    def plot_harmonics(self, t_range: Tuple[float, float] = (0, 10),
                      figsize: Tuple[int, int] = (12, 8),
                      x_for_dependencies: Optional[np.ndarray] = None) -> plt.Figure:
        """
        Строит графики гармоник с учетом возможных зависимостей от параметров
        """
        if not self.model.harmonics:
            raise ValueError("Нет гармонических колебаний")
        
        t = np.linspace(t_range[0], t_range[1], 1000)
        
        # Подготавливаем контекст для зависимых гармоник
        if x_for_dependencies is not None:
            context = self.model.evaluate_parameters(x_for_dependencies)
        else:
            # Если не задан x, используем значения по умолчанию
            context = {}
            for name, param in self.model.parameters.items():
                if isinstance(param, IndependentParameter):
                    x_default = np.linspace(param.x_range[0], param.x_range[1], 1)
                    context[name] = param.evaluate({'x': x_default})[0]
        
        n_plots = len(self.model.harmonics) + 1
        fig, axes = plt.subplots(n_plots, 1, figsize=(figsize[0], figsize[1] * n_plots/3))
        
        # Индивидуальные гармоники
        for i, h in enumerate(self.model.harmonics):
            harmonic_values = h.evaluate(t, context)
            axes[i].plot(t, harmonic_values, linewidth=1.5)
            axes[i].grid(True, alpha=0.3)
            
            # Показываем зависимости в заголовке
            deps = h.get_dependencies()
            deps_str = f" (зависит от: {', '.join(deps)})" if deps else ""
            axes[i].set_title(f"Гармоника {i+1}: {h}{deps_str}")
            axes[i].set_ylabel("Амплитуда")
        
        # Сумма гармоник
        total = np.zeros_like(t)
        for h in self.model.harmonics:
            total += h.evaluate(t, context)
        
        axes[-1].plot(t, total, 'r-', linewidth=2)
        axes[-1].grid(True, alpha=0.3)
        axes[-1].set_title("Сумма всех гармоник")
        axes[-1].set_xlabel("Время t")
        axes[-1].set_ylabel("Амплитуда")
        
        plt.tight_layout()
        self.figures["harmonics"] = fig
        return fig
    
    def plot_dependency_graph(self, figsize: Tuple[int, int] = (12, 8)) -> plt.Figure:
        """
        Визуализирует граф зависимостей между параметрами
        """
        try:
            import networkx as nx
        except ImportError:
            print("Для визуализации графа зависимостей установите networkx: pip install networkx")
            return None
        
        G = nx.DiGraph()
        
        # Добавляем узлы и ребра
        for name, param in self.model.parameters.items():
            node_attrs = {
                'color': 'lightblue' if isinstance(param, IndependentParameter) else 'lightgreen'
            }
            G.add_node(name, **node_attrs)
            
            if isinstance(param, DependentParameter):
                for dep in param.dependencies:
                    G.add_edge(dep, name)
        
        # Добавляем гармоники как узлы, если у них есть зависимости
        for i, h in enumerate(self.model.harmonics):
            deps = h.get_dependencies()
            if deps:
                harmonic_name = f"Harmonic_{i+1}"
                G.add_node(harmonic_name, color='lightcoral')
                for dep in deps:
                    G.add_edge(dep, harmonic_name)
        
        fig, ax = plt.subplots(figsize=figsize)
        
        # Раскладка графа
        pos = nx.spring_layout(G, k=2, iterations=50)
        
        # Рисуем граф
        node_colors = [G.nodes[node].get('color', 'lightgray') for node in G.nodes]
        nx.draw(G, pos, with_labels=True, node_color=node_colors,
                node_size=2000, font_size=10, font_weight='bold',
                arrows=True, arrowsize=20, ax=ax)
        
        ax.set_title("Граф зависимостей параметров", fontsize=16)
        
        self.figures["dependency_graph"] = fig
        return fig

# ====================== ПРИМЕР ИСПОЛЬЗОВАНИЯ ======================

def main():
    """Пример использования с зависимыми параметрами"""
    
    # Создаем модель
    model = MathematicalModel()
    # Константы
    U0 = 1.0
    h = 1.0
    pi = np.pi
    h_br = h / 2 / pi
    
    # Добавляем независимые параметры (базовые)
    model.add_independent_parameter(
        name="n",
        formula=lambda x: 1.0 * np.ones_like(x),  # Константа
        x_range=(0, 10),
        color='gray',
        description="Для дисперсионного уравнения"
    )
    
    model.add_independent_parameter(
        name="x0_br",
        formula=lambda x: 1.0 * np.ones_like(x),  # Константа
        x_range=(0, 10),
        color='blue',
        description="Порог"
    )
    
    model.add_independent_parameter(
        name="f0",
        formula=lambda x: 1.0 * np.ones_like(x),  # Константа
        x_range=(0, 10),
        color='red',
        description="Базовая частота"
    )
    
    # Добавляем зависимые параметры (выражаются через другие)
    model.add_dependent_parameter(
        name="Kn",
        expression=lambda A0, time: A0 * np.exp(-0.1 * time),  # Затухающая амплитуда
        dependencies=["n", "x0_br"],
        color='green',
        description="Эффективная амплитуда (зависит от времени)"
    )
    
    model.add_dependent_parameter(
        name="f_effective",
        expression=lambda f0, time: f0 * (1 + 0.05 * np.sin(time)),  # Модулированная частота
        dependencies=["f0", "time"],
        color='orange',
        description="Эффективная частота (модулирована)"
    )
    
    model.add_dependent_parameter(
        name="modulation_index",
        expression=lambda A_effective, f_effective: A_effective * f_effective / 2,
        dependencies=["A_effective", "f_effective"],
        color='purple',
        description="Индекс модуляции (произведение)"
    )
    
    # Добавляем гармоники, которые могут зависеть от параметров
    model.add_harmonic(
        amplitude="A_effective",  # Использует параметр A_effective
        frequency="f_effective",  # Использует параметр f_effective
        phase=0,
        amplitude_depends_on=["A_effective"],
        frequency_depends_on=["f_effective"]
    )
    
    model.add_harmonic(
        amplitude=0.5,
        frequency=2.0,
        phase=np.pi/4
    )
    
    model.add_harmonic(
        amplitude="modulation_index",
        frequency=3.0,
        phase=0,
        amplitude_depends_on=["modulation_index"]
    )
    
    # Устанавливаем основную формулу
    def main_formula(x, A_effective, f_effective, modulation_index, harmonic_sum):
        """
        Основная формула: комбинация параметров и гармоник
        """
        # Параметрическая часть
        parametric_part = A_effective * np.sin(2 * np.pi * f_effective * x)
        
        # Модуляционная часть
        modulation_part = modulation_index * np.cos(2 * np.pi * x)
        
        # Гармоническая часть
        harmonic_part = harmonic_sum(x)
        
        return parametric_part + modulation_part + harmonic_part
    
    model.set_main_formula(main_formula)
    
    # Создаем визуализатор
    viz = ModelVisualizer(model)
    
    # Строим граф зависимостей
    print("Визуализация графа зависимостей...")
    viz.plot_dependency_graph()
    
    # Строим графики параметров с зависимостями
    print("\nГрафики параметров с зависимостями...")
    viz.plot_parameter("A_effective", show_dependencies=True)
    viz.plot_parameter("f_effective", show_dependencies=True)
    viz.plot_parameter_with_dependencies("modulation_index")
    
    # Строим графики всех параметров
    print("\nВсе параметры...")
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()
    
    x_test = np.linspace(0, 10, 1000)
    context = model.evaluate_parameters(x_test)
    
    for i, (name, param) in enumerate(model.parameters.items()):
        if i < len(axes):
            axes[i].plot(x_test, context[name], color=param.color, 
                        linestyle=param.line_style, linewidth=2)
            axes[i].grid(True, alpha=0.3)
            axes[i].set_title(f"{name}\n{param.description}")
            axes[i].set_xlabel("x")
    
    # Скрываем лишний подграфик
    for j in range(len(model.parameters), len(axes)):
        axes[j].set_visible(False)
    
    plt.tight_layout()
    
    # Строим графики гармоник с зависимостями
    print("\nГармонические колебания...")
    viz.plot_harmonics(t_range=(0, 5), x_for_dependencies=x_test)
    
    # Строим график основной формулы
    print("\nОсновная формула...")
    fig, ax = plt.subplots(figsize=(12, 6))
    y_main = model.evaluate_main(x_test)
    ax.plot(x_test, y_main, 'b-', linewidth=2)
    ax.grid(True, alpha=0.3)
    ax.set_title("Результат работы основной формулы", fontsize=14)
    ax.set_xlabel("x", fontsize=12)
    ax.set_ylabel("F(x)", fontsize=12)
    
    # Показываем все графики
    plt.show()
    
    # Выводим информацию о зависимостях
    print("\nСтруктура зависимостей:")
    deps = model.get_all_dependencies()
    for param, dependencies in deps.items():
        print(f"  {param} зависит от: {', '.join(dependencies)}")

if __name__ == "__main__":
    main()