"""Feasibility-First Search Framework (P3-M2).

Pure Python implementation (no numpy/scipy/sklearn dependency).
Implements:
- LHS (Latin Hypercube Sampling) for initial space coverage
- Active learning batch selection (distance + uncertainty heuristic)
- Constrained feasibility prioritization
- Local trust region refinement
- Search state management with checkpoint support

Per third-party review: default mode is feasibility-first
(constrained Bayesian / active learning), not fixed full-factorial DoE.
"""
import math
import random
import hashlib
import json
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any, Callable
from datetime import datetime

from .l0_prescreening import L0PreScreeningEngine, FeasibilityReport


@dataclass
class ParameterRange:
    """Definition of a scan parameter range."""
    name: str
    min_value: float
    max_value: float
    step: Optional[float] = None  # if set, values are quantized
    unit: str = ""
    description: str = ""

    def normalize(self, value: float) -> float:
        """Normalize value to [0, 1] range."""
        if self.max_value == self.min_value:
            return 0.5
        return (value - self.min_value) / (self.max_value - self.min_value)

    def denormalize(self, norm: float) -> float:
        """Denormalize from [0, 1] to actual range."""
        value = self.min_value + norm * (self.max_value - self.min_value)
        if self.step is not None and self.step > 0:
            value = round(value / self.step) * self.step
            value = max(self.min_value, min(self.max_value, value))
        return value

    def random_value(self, rng: random.Random) -> float:
        """Generate a random value within range."""
        return self.denormalize(rng.random())


@dataclass
class SearchPoint:
    """A single point in the search space."""
    id: int
    params: Dict[str, float]
    status: str = "pending"  # pending / running / ok / failed / infeasible
    metrics: Dict[str, float] = field(default_factory=dict)
    feasibility_report: Optional[Dict[str, Any]] = None
    surrogate_prediction: Optional[float] = None
    surrogate_uncertainty: Optional[float] = None
    batch_id: int = 0
    created_at: str = ""

    def param_hash(self) -> str:
        """Generate a hash of parameter values for caching/dedup."""
        sorted_params = sorted(self.params.items())
        param_str = json.dumps(sorted_params, sort_keys=True)
        return hashlib.md5(param_str.encode()).hexdigest()[:12]


@dataclass
class SearchState:
    """Complete state of an adaptive search run."""
    run_id: str
    parameters: List[ParameterRange]
    points: List[SearchPoint] = field(default_factory=list)
    current_batch: int = 0
    total_budget: int = 80
    used_budget: int = 0
    batch_size: int = 4
    search_method: str = "active_learning"
    trust_region_active: bool = False
    trust_region_center: Optional[Dict[str, float]] = None
    trust_region_radius: float = 0.3  # fraction of normalized range
    best_feasible_point: Optional[SearchPoint] = None
    best_objective_value: float = float('inf')
    objective_metric: str = "tavg_nm"  # metric to optimize
    objective_direction: str = "maximize"  # maximize / minimize
    convergence_status: str = "searching"  # searching / converged / stalled / budget_exhausted
    history: List[Dict[str, Any]] = field(default_factory=list)
    created_at: str = ""
    updated_at: str = ""

    def get_param_names(self) -> List[str]:
        return [p.name for p in self.parameters]

    def get_param_by_name(self, name: str) -> Optional[ParameterRange]:
        for p in self.parameters:
            if p.name == name:
                return p
        return None

    def get_pending_points(self) -> List[SearchPoint]:
        return [p for p in self.points if p.status == "pending"]

    def get_completed_points(self) -> List[SearchPoint]:
        return [p for p in self.points if p.status in ("ok", "failed", "infeasible")]

    def get_feasible_points(self) -> List[SearchPoint]:
        return [p for p in self.points if p.status == "ok" and p.feasibility_report and p.feasibility_report.get("feasible", False)]

    def remaining_budget(self) -> int:
        return self.total_budget - self.used_budget


class FeasibilityFirstSearch:
    """Feasibility-first adaptive search engine.

    Implements the recommended default path from third-party review:
    L0 pre-screening + initial samples + active learning batch selection
    + local trust region refinement.
    """

    def __init__(
        self,
        parameters: List[ParameterRange],
        l0_engine: Optional[L0PreScreeningEngine] = None,
        total_budget: int = 80,
        batch_size: int = 4,
        initial_samples: int = 16,
        objective_metric: str = "tavg_nm",
        objective_direction: str = "maximize",
        seed: int = 42,
    ):
        self.rng = random.Random(seed)
        self.parameters = parameters
        self.l0_engine = l0_engine or L0PreScreeningEngine()
        self.total_budget = total_budget
        self.batch_size = batch_size
        self.initial_samples = initial_samples
        self.objective_metric = objective_metric
        self.objective_direction = objective_direction
        self.state = SearchState(
            run_id=f"search_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            parameters=parameters,
            total_budget=total_budget,
            batch_size=batch_size,
            objective_metric=objective_metric,
            objective_direction=objective_direction,
            created_at=datetime.now().isoformat(),
        )

    def latin_hypercube_sample(self, n_samples: int) -> List[Dict[str, float]]:
        """Generate Latin Hypercube Samples in normalized space.

        Pure Python implementation (no numpy dependency).
        """
        n_params = len(self.parameters)
        # Create n_samples intervals per dimension
        intervals = []
        for _ in range(n_params):
            # Shuffle interval assignments
            assignments = list(range(n_samples))
            self.rng.shuffle(assignments)
            intervals.append(assignments)

        samples = []
        for i in range(n_samples):
            sample = {}
            for j, param in enumerate(self.parameters):
                # Random point within assigned interval
                interval_idx = intervals[j][i]
                norm = (interval_idx + self.rng.random()) / n_samples
                sample[param.name] = param.denormalize(norm)
            samples.append(sample)
        return samples

    def _min_distance_to_existing(self, params: Dict[str, float], existing: List[Dict[str, float]]) -> float:
        """Calculate minimum normalized Euclidean distance to existing points."""
        if not existing:
            return float('inf')
        min_dist = float('inf')
        for ex in existing:
            dist_sq = 0.0
            for param in self.parameters:
                v1 = param.normalize(params.get(param.name, param.min_value))
                v2 = param.normalize(ex.get(param.name, param.min_value))
                dist_sq += (v1 - v2) ** 2
            dist = math.sqrt(dist_sq)
            min_dist = min(min_dist, dist)
        return min_dist

    def _is_duplicate(self, params: Dict[str, float], existing_points: List[SearchPoint]) -> bool:
        """Check if parameter set is a duplicate of an existing point."""
        param_hash = hashlib.md5(
            json.dumps(sorted(params.items()), sort_keys=True).encode()
        ).hexdigest()[:12]
        return any(p.param_hash() == param_hash for p in existing_points)

    def generate_initial_batch(self) -> List[SearchPoint]:
        """Generate initial LHS samples, filtered by L0 feasibility.

        Prioritizes feasible points but includes some infeasible for
        boundary learning (active learning needs both classes).
        """
        # Generate more samples than needed, then select diverse subset
        n_candidates = max(self.initial_samples * 3, 20)
        candidates = self.latin_hypercube_sample(n_candidates)

        # L0 pre-screen all candidates
        scored = []
        for params in candidates:
            report = self.l0_engine.evaluate(params)
            scored.append((params, report))

        # Separate feasible and infeasible
        feasible = [(p, r) for p, r in scored if r.feasible]
        infeasible = [(p, r) for p, r in scored if not r.feasible]

        # Select diverse subset: mostly feasible, some infeasible for boundary
        n_feasible = min(len(feasible), int(self.initial_samples * 0.75))
        n_infeasible = min(len(infeasible), self.initial_samples - n_feasible)

        selected = []
        selected_params = []

        # Greedy max-distance selection for feasible points
        feasible.sort(key=lambda x: x[1].pass_rate, reverse=True)
        for params, report in feasible:
            if len(selected) >= n_feasible:
                break
            # Check duplicate against already selected params
            is_dup = any(
                hashlib.md5(json.dumps(sorted(p.items()), sort_keys=True).encode()).hexdigest()[:12] ==
                hashlib.md5(json.dumps(sorted(params.items()), sort_keys=True).encode()).hexdigest()[:12]
                for p, _ in selected
            )
            if not is_dup:
                dist = self._min_distance_to_existing(params, selected_params)
                if dist > 0.1 or len(selected) < 3:
                    selected.append((params, report))
                    selected_params.append(params)

        # Add infeasible points for boundary learning
        for params, report in infeasible:
            if len(selected) >= self.initial_samples:
                break
            is_dup = any(
                hashlib.md5(json.dumps(sorted(p.items()), sort_keys=True).encode()).hexdigest()[:12] ==
                hashlib.md5(json.dumps(sorted(params.items()), sort_keys=True).encode()).hexdigest()[:12]
                for p, _ in selected
            )
            if not is_dup:
                selected.append((params, report))
                selected_params.append(params)

        # Create SearchPoint objects
        points = []
        for i, (params, report) in enumerate(selected):
            point = SearchPoint(
                id=len(self.state.points) + i,
                params=params,
                status="infeasible" if not report.feasible else "pending",
                feasibility_report=report.to_dict(),
                batch_id=0,
                created_at=datetime.now().isoformat(),
            )
            if not report.feasible:
                # Mark as completed (L0 rejected) without using solver budget
                point.status = "infeasible"
            points.append(point)

        self.state.points.extend(points)
        self.state.current_batch = 1
        self.state.used_budget += sum(1 for p in points if p.status == "pending")
        self.state.updated_at = datetime.now().isoformat()

        return [p for p in points if p.status == "pending"]

    def select_next_batch(self) -> List[SearchPoint]:
        """Select next batch of points using active learning.

        Strategy:
        1. If feasible points found, activate trust region around best
        2. Select points balancing:
           - Exploration: high distance from existing points
           - Exploitation: near best feasible point (trust region)
           - Uncertainty: points near feasibility boundary
        3. L0 pre-screen all candidates, reject obviously infeasible
        """
        if self.state.remaining_budget() <= 0:
            self.state.convergence_status = "budget_exhausted"
            return []

        feasible_points = self.state.get_feasible_points()
        completed_params = [p.params for p in self.state.get_completed_points()]

        # Activate trust region if we have enough feasible points
        if len(feasible_points) >= 3 and not self.state.trust_region_active:
            # C4 fix: select anchor by optimization direction
            if self.objective_direction == "minimize":
                best = min(feasible_points, key=lambda p: p.metrics.get(self.objective_metric, float("inf")))
            else:
                best = max(feasible_points, key=lambda p: p.metrics.get(self.objective_metric, 0))
            self.state.trust_region_active = True
            self.state.trust_region_center = best.params
            self.state.best_feasible_point = best
            self.state.history.append({
                "event": "trust_region_activated",
                "center": best.params,
                "batch": self.state.current_batch,
            })

        # Generate candidates
        n_candidates = 50
        candidates = []

        if self.state.trust_region_active and self.state.trust_region_center:
            # 70% from trust region, 30% from global exploration
            n_trust = int(n_candidates * 0.7)
            n_global = n_candidates - n_trust

            # Trust region samples (Gaussian-like around center)
            for _ in range(n_trust):
                params = {}
                for param in self.parameters:
                    center_norm = param.normalize(self.state.trust_region_center[param.name])
                    # Sample with decreasing radius as search progresses
                    radius = self.state.trust_region_radius * max(0.3, 1.0 - self.state.used_budget / self.total_budget)
                    sample_norm = center_norm + self.rng.gauss(0, radius * 0.3)
                    sample_norm = max(0.0, min(1.0, sample_norm))
                    params[param.name] = param.denormalize(sample_norm)
                candidates.append(params)

            # Global exploration samples
            candidates.extend(self.latin_hypercube_sample(n_global))
        else:
            # No trust region yet: full LHS exploration
            candidates = self.latin_hypercube_sample(n_candidates)

        # Score and select candidates
        scored_candidates = []
        for params in candidates:
            if self._is_duplicate(params, self.state.points):
                continue

            # L0 pre-screen
            report = self.l0_engine.evaluate(params)
            if not report.feasible:
                continue  # Skip obviously infeasible

            # Distance score (exploration)
            dist = self._min_distance_to_existing(params, completed_params)
            dist_score = min(dist / math.sqrt(len(self.parameters)), 1.0)

            # Trust region proximity score (exploitation)
            trust_score = 0.0
            if self.state.trust_region_center:
                dist_to_center_sq = 0.0
                for param in self.parameters:
                    v1 = param.normalize(params[param.name])
                    v2 = param.normalize(self.state.trust_region_center[param.name])
                    dist_to_center_sq += (v1 - v2) ** 2
                dist_to_center = math.sqrt(dist_to_center_sq)
                trust_score = max(0, 1.0 - dist_to_center / self.state.trust_region_radius)

            # Combined score (balance exploration and exploitation)
            if self.state.trust_region_active:
                score = 0.4 * dist_score + 0.6 * trust_score
            else:
                score = dist_score

            scored_candidates.append((params, report, score))

        # Sort by score and select top batch_size
        scored_candidates.sort(key=lambda x: x[2], reverse=True)
        batch_size = min(self.batch_size, self.state.remaining_budget(), len(scored_candidates))

        selected = []
        for i in range(batch_size):
            params, report, score = scored_candidates[i]
            point = SearchPoint(
                id=len(self.state.points) + i,
                params=params,
                status="pending",
                feasibility_report=report.to_dict(),
                batch_id=self.state.current_batch,
                created_at=datetime.now().isoformat(),
            )
            selected.append(point)

        self.state.points.extend(selected)
        self.state.used_budget += len(selected)
        self.state.current_batch += 1
        self.state.updated_at = datetime.now().isoformat()

        return selected

    def report_result(self, point_id: int, metrics: Dict[str, float], status: str = "ok") -> None:
        """Report simulation result for a point.

        Args:
            point_id: ID of the search point
            metrics: Dictionary of metric values
            status: ok / failed
        """
        for point in self.state.points:
            if point.id == point_id:
                point.metrics = metrics
                point.status = status

                # Update best feasible point
                if status == "ok" and self.objective_metric in metrics:
                    value = metrics[self.objective_metric]
                    is_better = (
                        (self.objective_direction == "maximize" and value > self.state.best_objective_value) or
                        (self.objective_direction == "minimize" and value < self.state.best_objective_value)
                    )
                    if is_better or self.state.best_feasible_point is None:
                        self.state.best_objective_value = value
                        self.state.best_feasible_point = point
                        # C5 fix: migrate trust region center to new best point
                        if self.state.trust_region_active:
                            self.state.trust_region_center = point.params

                # Check convergence
                self._check_convergence()
                break

        self.state.updated_at = datetime.now().isoformat()

    def _check_convergence(self) -> None:
        """Check if search has converged or stalled."""
        feasible = self.state.get_feasible_points()
        if len(feasible) < 5:
            return

        # Check if objective has improved in last N points
        recent = feasible[-10:] if len(feasible) >= 10 else feasible
        if len(recent) >= 5:
            values = [p.metrics.get(self.objective_metric, 0) for p in recent]
            if self.objective_direction == "maximize":
                improvement = max(values) - max(values[:-3]) if len(values) > 3 else 0
            else:
                improvement = min(values[:-3]) - min(values) if len(values) > 3 else 0

            if improvement < 0.001 and self.state.trust_region_active:
                # Shrink trust region
                self.state.trust_region_radius *= 0.7
                if self.state.trust_region_radius < 0.05:
                    self.state.convergence_status = "converged"
                    self.state.history.append({
                        "event": "converged",
                        "reason": "trust_region_shrunk_below_threshold",
                        "batch": self.state.current_batch,
                    })

    def get_state_summary(self) -> Dict[str, Any]:
        """Get a summary of current search state."""
        return {
            "run_id": self.state.run_id,
            "search_method": self.state.search_method,
            "convergence_status": self.state.convergence_status,
            "total_budget": self.state.total_budget,
            "used_budget": self.state.used_budget,
            "remaining_budget": self.state.remaining_budget(),
            "current_batch": self.state.current_batch,
            "total_points": len(self.state.points),
            "pending_points": len(self.state.get_pending_points()),
            "completed_points": len(self.state.get_completed_points()),
            "feasible_points": len(self.state.get_feasible_points()),
            "trust_region_active": self.state.trust_region_active,
            "trust_region_radius": round(self.state.trust_region_radius, 4),
            "best_objective_value": self.state.best_objective_value if self.state.best_objective_value != float('inf') else None,
            "best_point_params": self.state.best_feasible_point.params if self.state.best_feasible_point else None,
        }

    def export_state(self) -> Dict[str, Any]:
        """Export full search state for checkpointing."""
        return {
            "state": {
                "run_id": self.state.run_id,
                "current_batch": self.state.current_batch,
                "total_budget": self.state.total_budget,
                "used_budget": self.state.used_budget,
                "convergence_status": self.state.convergence_status,
                "trust_region_active": self.state.trust_region_active,
                "trust_region_center": self.state.trust_region_center,
                "trust_region_radius": self.state.trust_region_radius,
                "best_objective_value": self.state.best_objective_value,
            },
            "parameters": [
                {"name": p.name, "min": p.min_value, "max": p.max_value, "step": p.step, "unit": p.unit}
                for p in self.parameters
            ],
            "points": [
                {
                    "id": p.id,
                    "params": p.params,
                    "status": p.status,
                    "metrics": p.metrics,
                    "batch_id": p.batch_id,
                    "feasible": p.feasibility_report.get("feasible") if p.feasibility_report else None,
                }
                for p in self.state.points
            ],
        }