"""
Design Space Insight

Provides interactive visualization of the design space explored during optimization.
Shows parameter relationships, objective landscapes, and design evolution.

This insight bridges optimization metrics (from Analysis) with physics understanding,
showing how design parameters affect the physical objectives.

Applicable to: All optimization studies with completed trials.
"""

from pathlib import Path
from datetime import datetime
from typing import Dict, Any, List, Optional, Tuple
import sqlite3
import json
import numpy as np

from .base import StudyInsight, InsightConfig, InsightResult, register_insight

# Lazy imports
_plotly_loaded = False
_go = None
_make_subplots = None


def _load_dependencies():
    """Lazy load heavy dependencies."""
    global _plotly_loaded, _go, _make_subplots
    if not _plotly_loaded:
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots
        _go = go
        _make_subplots = make_subplots
        _plotly_loaded = True


@register_insight
class DesignSpaceInsight(StudyInsight):
    """
    Design space exploration visualization.

    Shows:
    - Parallel coordinates plot of parameters vs objectives
    - Scatter matrix of parameter relationships
    - 3D parameter-objective landscape
    - Best design summary with physics interpretation
    """

    insight_type = "design_space"
    name = "Design Space Explorer"
    description = "Interactive parameter-objective relationship visualization"
    category = "design_exploration"
    applicable_to = ["all"]  # Works with any optimization study
    required_files = []  # Requires study.db, not OP2

    def __init__(self, study_path: Path):
        super().__init__(study_path)
        self.db_path = self.results_path / "study.db"
        self._trials: Optional[List[Dict]] = None
        self._params: Optional[List[str]] = None
        self._objectives: Optional[List[str]] = None

    def can_generate(self) -> bool:
        """Check if study.db exists with trial data."""
        if not self.db_path.exists():
            return False

        try:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'")
            count = cursor.fetchone()[0]
            conn.close()
            return count >= 5  # Need at least 5 trials
        except Exception:
            return False

    def _load_data(self):
        """Load trial data from study.db."""
        if self._trials is not None:
            return

        conn = sqlite3.connect(str(self.db_path))
        cursor = conn.cursor()

        # Get completed trials
        cursor.execute("""
            SELECT trial_id, params, values, state
            FROM trials
            WHERE state = 'COMPLETE'
            ORDER BY trial_id
        """)

        self._trials = []
        self._params = None
        self._objectives = None

        for row in cursor.fetchall():
            trial_id, params_json, values_json, state = row

            try:
                params = json.loads(params_json) if params_json else {}
                values = json.loads(values_json) if values_json else {}
            except json.JSONDecodeError:
                continue

            # Flatten nested values
            flat_values = {}
            for k, v in values.items():
                if isinstance(v, dict):
                    flat_values.update(v)
                else:
                    flat_values[k] = v

            self._trials.append({
                'trial_id': trial_id,
                'params': params,
                'values': flat_values,
            })

            # Extract param and objective names from first trial
            if self._params is None:
                self._params = list(params.keys())
            if self._objectives is None:
                self._objectives = list(flat_values.keys())

        conn.close()

    def _generate(self, config: InsightConfig) -> InsightResult:
        """Generate design space visualization."""
        self._load_data()

        if not self._trials or len(self._trials) < 5:
            return InsightResult(success=False,
                                 error=f"Need at least 5 trials, found: {len(self._trials or [])}")

        _load_dependencies()

        # Configuration
        colorscale = config.extra.get('colorscale', 'Viridis')
        primary_objective = config.extra.get('primary_objective', None)

        # Use first objective if not specified
        if primary_objective is None and self._objectives:
            primary_objective = self._objectives[0]

        # Build data arrays
        n_trials = len(self._trials)
        param_data = {p: [] for p in self._params}
        obj_data = {o: [] for o in self._objectives}
        trial_ids = []

        for trial in self._trials:
            trial_ids.append(trial['trial_id'])
            for p in self._params:
                param_data[p].append(trial['params'].get(p, np.nan))
            for o in self._objectives:
                obj_data[o].append(trial['values'].get(o, np.nan))

        # Convert to arrays
        for p in self._params:
            param_data[p] = np.array(param_data[p])
        for o in self._objectives:
            obj_data[o] = np.array(obj_data[o])

        # Find best trial
        if primary_objective and primary_objective in obj_data:
            obj_values = obj_data[primary_objective]
            valid_mask = ~np.isnan(obj_values)
            if np.any(valid_mask):
                best_idx = np.nanargmin(obj_values)
                best_trial = self._trials[best_idx]
                best_value = obj_values[best_idx]
            else:
                best_trial = None
                best_value = None
        else:
            best_trial = None
            best_value = None

        # Build visualization
        n_params = len(self._params)
        n_objs = len(self._objectives)

        # Layout: 2x2 grid
        fig = _make_subplots(
            rows=2, cols=2,
            specs=[
                [{"type": "parcoords", "colspan": 2}, None],
                [{"type": "scatter3d" if n_params >= 2 else "xy"},
                 {"type": "table"}]
            ],
            row_heights=[0.55, 0.45],
            subplot_titles=[
                "<b>Parallel Coordinates - Design Space</b>",
                "<b>Parameter Landscape</b>",
                "<b>Best Design</b>"
            ]
        )

        # 1. Parallel coordinates
        dimensions = []

        # Add parameters
        for p in self._params:
            values = param_data[p]
            if not np.all(np.isnan(values)):
                dimensions.append(dict(
                    label=p,
                    values=values,
                    range=[float(np.nanmin(values)), float(np.nanmax(values))]
                ))

        # Add objectives
        for o in self._objectives:
            values = obj_data[o]
            if not np.all(np.isnan(values)):
                dimensions.append(dict(
                    label=o,
                    values=values,
                    range=[float(np.nanmin(values)), float(np.nanmax(values))]
                ))

        if dimensions:
            # Color by primary objective
            color_values = obj_data.get(primary_objective, trial_ids)
            if isinstance(color_values, list):
                color_values = np.array(color_values)

            fig.add_trace(_go.Parcoords(
                line=dict(
                    color=color_values,
                    colorscale=colorscale,
                    showscale=True,
                    colorbar=dict(title=primary_objective or "Trial", thickness=15)
                ),
                dimensions=dimensions,
            ), row=1, col=1)

        # 2. 3D Parameter landscape (first 2 params vs primary objective)
        if n_params >= 2 and primary_objective:
            x_param = self._params[0]
            y_param = self._params[1]
            z_values = obj_data.get(primary_objective, [])

            fig.add_trace(_go.Scatter3d(
                x=param_data[x_param],
                y=param_data[y_param],
                z=z_values,
                mode='markers',
                marker=dict(
                    size=6,
                    color=z_values,
                    colorscale=colorscale,
                    opacity=0.8,
                    showscale=False,
                ),
                text=[f"Trial {tid}" for tid in trial_ids],
                hovertemplate=(
                    f"{x_param}: %{{x:.3f}}<br>"
                    f"{y_param}: %{{y:.3f}}<br>"
                    f"{primary_objective}: %{{z:.4f}}<br>"
                    "%{text}<extra></extra>"
                ),
            ), row=2, col=1)

            # Highlight best point
            if best_trial:
                fig.add_trace(_go.Scatter3d(
                    x=[best_trial['params'].get(x_param)],
                    y=[best_trial['params'].get(y_param)],
                    z=[best_value],
                    mode='markers',
                    marker=dict(size=12, color='red', symbol='diamond'),
                    name='Best',
                    showlegend=True,
                ), row=2, col=1)

            fig.update_scenes(
                xaxis_title=x_param,
                yaxis_title=y_param,
                zaxis_title=primary_objective,
            )
        elif n_params >= 1 and primary_objective:
            # 2D scatter
            x_param = self._params[0]
            z_values = obj_data.get(primary_objective, [])

            fig.add_trace(_go.Scatter(
                x=param_data[x_param],
                y=z_values,
                mode='markers',
                marker=dict(size=8, color=z_values, colorscale=colorscale),
            ), row=2, col=1)

            fig.update_xaxes(title_text=x_param, row=2, col=1)
            fig.update_yaxes(title_text=primary_objective, row=2, col=1)

        # 3. Best design table
        if best_trial:
            labels = ["<b>Metric</b>", "<b>Value</b>"]

            # Combine params and objectives
            table_labels = ["Trial ID"] + self._params + self._objectives
            table_values = [str(best_trial['trial_id'])]

            for p in self._params:
                val = best_trial['params'].get(p, 'N/A')
                table_values.append(f"{val:.4f}" if isinstance(val, (int, float)) else str(val))

            for o in self._objectives:
                val = best_trial['values'].get(o, 'N/A')
                table_values.append(f"{val:.4f}" if isinstance(val, (int, float)) else str(val))

            fig.add_trace(_go.Table(
                header=dict(values=labels,
                            fill_color='#1f2937', font=dict(color='white')),
                cells=dict(values=[table_labels, table_values],
                           fill_color='#374151', font=dict(color='white'))
            ), row=2, col=2)
        else:
            fig.add_trace(_go.Table(
                header=dict(values=["<b>Info</b>"],
                            fill_color='#1f2937', font=dict(color='white')),
                cells=dict(values=[["No valid trials found"]],
                           fill_color='#374151', font=dict(color='white'))
            ), row=2, col=2)

        # Layout
        fig.update_layout(
            width=1500, height=1000,
            paper_bgcolor='#111827', plot_bgcolor='#1f2937',
            font=dict(color='white'),
            title=dict(
                text=f"<b>Atomizer Design Space Explorer</b><br>"
                     f"<sub>{n_trials} trials, {n_params} parameters, {n_objs} objectives</sub>",
                x=0.5, font=dict(size=18)
            ),
        )

        # Save HTML
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = config.output_dir or self.insights_path
        output_dir.mkdir(parents=True, exist_ok=True)

        html_path = output_dir / f"design_space_{timestamp}.html"
        html_path.write_text(
            fig.to_html(include_plotlyjs='cdn', full_html=True),
            encoding='utf-8'
        )

        # Summary
        summary = {
            'n_trials': n_trials,
            'n_params': n_params,
            'n_objectives': n_objs,
            'parameters': self._params,
            'objectives': self._objectives,
        }

        if best_trial:
            summary['best_trial_id'] = best_trial['trial_id']
            summary['best_params'] = best_trial['params']
            summary['best_values'] = best_trial['values']

        return InsightResult(
            success=True,
            html_path=html_path,
            plotly_figure=fig.to_dict(),
            summary=summary
        )