"""
Experiment tracking for Neural models.

This module provides functionality to track experiments, including hyperparameters,
metrics, and artifacts.
"""

import os
import json
import logging
import time
import datetime
import uuid
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Any, Optional, Union, Tuple
from pathlib import Path

# Configure logger
logger = logging.getLogger(__name__)

class ExperimentTracker:
    """Tracks experiments, including hyperparameters, metrics, and artifacts."""
    
    def __init__(self, experiment_name: str = None, base_dir: str = "neural_experiments"):
        """
        Initialize the experiment tracker.
        
        Args:
            experiment_name: Name of the experiment (defaults to timestamp if None)
            base_dir: Base directory for storing experiment data
        """
        self.experiment_name = experiment_name or f"experiment_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.experiment_id = str(uuid.uuid4())[:8]
        self.base_dir = base_dir
        self.experiment_dir = os.path.join(base_dir, f"{self.experiment_name}_{self.experiment_id}")
        self.metrics_history = []
        self.hyperparameters = {}
        self.metadata = {
            "start_time": datetime.datetime.now().isoformat(),
            "status": "created"
        }
        self.artifacts = {}
        
        # Create experiment directory
        os.makedirs(self.experiment_dir, exist_ok=True)
        os.makedirs(os.path.join(self.experiment_dir, "artifacts"), exist_ok=True)
        os.makedirs(os.path.join(self.experiment_dir, "plots"), exist_ok=True)
        
        # Save initial metadata
        self._save_metadata()
        
        logger.info(f"Initialized experiment tracker: {self.experiment_name} (ID: {self.experiment_id})")
    
    def log_hyperparameters(self, hyperparameters: Dict[str, Any]):
        """
        Log hyperparameters for the experiment.
        
        Args:
            hyperparameters: Dictionary of hyperparameters
        """
        self.hyperparameters.update(hyperparameters)
        self._save_hyperparameters()
        logger.debug(f"Logged hyperparameters: {hyperparameters}")
    
    def log_metrics(self, metrics: Dict[str, float], step: int = None):
        """
        Log metrics for the experiment.
        
        Args:
            metrics: Dictionary of metric names and values
            step: Step or epoch number (optional)
        """
        timestamp = time.time()
        metrics_entry = {
            "timestamp": timestamp,
            "step": step,
            **metrics
        }
        self.metrics_history.append(metrics_entry)
        self._save_metrics()
        logger.debug(f"Logged metrics at step {step}: {metrics}")
    
    def log_artifact(self, artifact_path: str, artifact_name: str = None):
        """
        Log an artifact for the experiment.
        
        Args:
            artifact_path: Path to the artifact file
            artifact_name: Name to use for the artifact (defaults to filename)
        """
        if not os.path.exists(artifact_path):
            logger.error(f"Artifact not found: {artifact_path}")
            return
            
        artifact_name = artifact_name or os.path.basename(artifact_path)
        artifact_dest = os.path.join(self.experiment_dir, "artifacts", artifact_name)
        
        # Copy the artifact
        import shutil
        shutil.copy2(artifact_path, artifact_dest)
        
        # Record the artifact
        self.artifacts[artifact_name] = {
            "path": artifact_dest,
            "type": self._get_artifact_type(artifact_path),
            "size": os.path.getsize(artifact_path),
            "timestamp": time.time()
        }
        
        self._save_artifacts()
        logger.debug(f"Logged artifact: {artifact_name}")
    
    def log_model(self, model_path: str, framework: str = "unknown"):
        """
        Log a model for the experiment.
        
        Args:
            model_path: Path to the model file
            framework: Framework used for the model (tensorflow, pytorch, etc.)
        """
        model_name = os.path.basename(model_path)
        model_dest = os.path.join(self.experiment_dir, "artifacts", model_name)
        
        # Copy the model
        import shutil
        shutil.copy2(model_path, model_dest)
        
        # Record the model
        self.artifacts[model_name] = {
            "path": model_dest,
            "type": "model",
            "framework": framework,
            "size": os.path.getsize(model_path),
            "timestamp": time.time()
        }
        
        self._save_artifacts()
        logger.debug(f"Logged {framework} model: {model_name}")
    
    def log_figure(self, figure: plt.Figure, figure_name: str):
        """
        Log a matplotlib figure for the experiment.
        
        Args:
            figure: Matplotlib figure
            figure_name: Name for the figure
        """
        if not figure_name.endswith(('.png', '.jpg', '.jpeg', '.svg')):
            figure_name += '.png'
            
        figure_path = os.path.join(self.experiment_dir, "plots", figure_name)
        figure.savefig(figure_path)
        
        # Record the figure
        self.artifacts[figure_name] = {
            "path": figure_path,
            "type": "figure",
            "size": os.path.getsize(figure_path),
            "timestamp": time.time()
        }
        
        self._save_artifacts()
        logger.debug(f"Logged figure: {figure_name}")
    
    def set_status(self, status: str):
        """
        Set the status of the experiment.
        
        Args:
            status: Status of the experiment (e.g., 'running', 'completed', 'failed')
        """
        self.metadata["status"] = status
        if status in ["completed", "failed"]:
            self.metadata["end_time"] = datetime.datetime.now().isoformat()
        self._save_metadata()
        logger.info(f"Set experiment status to: {status}")
    
    def get_metrics(self, metric_name: str = None) -> Union[List[Dict[str, Any]], List[float]]:
        """
        Get metrics history for the experiment.
        
        Args:
            metric_name: Name of the metric to retrieve (optional)
            
        Returns:
            List of metric values or full metrics history
        """
        if metric_name:
            return [entry.get(metric_name, None) for entry in self.metrics_history if metric_name in entry]
        return self.metrics_history
    
    def get_best_metric(self, metric_name: str, mode: str = "max") -> Tuple[float, int]:
        """
        Get the best value for a metric and its corresponding step.
        
        Args:
            metric_name: Name of the metric
            mode: 'max' or 'min' depending on whether higher or lower is better
            
        Returns:
            Tuple of (best_value, step)
        """
        if not self.metrics_history:
            return None, None
            
        metric_values = [(entry.get(metric_name, None), entry.get("step", i)) 
                         for i, entry in enumerate(self.metrics_history) 
                         if metric_name in entry]
        
        if not metric_values:
            return None, None
            
        if mode == "max":
            best_value, step = max(metric_values, key=lambda x: x[0] if x[0] is not None else float('-inf'))
        else:
            best_value, step = min(metric_values, key=lambda x: x[0] if x[0] is not None else float('inf'))
            
        return best_value, step
    
    def plot_metrics(self, metric_names: List[str] = None, figsize: Tuple[int, int] = (10, 6)) -> plt.Figure:
        """
        Plot metrics history.
        
        Args:
            metric_names: List of metric names to plot (plots all metrics if None)
            figsize: Figure size
            
        Returns:
            Matplotlib figure
        """
        if not self.metrics_history:
            fig, ax = plt.subplots(figsize=figsize)
            ax.text(0.5, 0.5, "No metrics data available", ha='center', va='center')
            return fig
            
        # Determine which metrics to plot
        if metric_names is None:
            # Get all metric names excluding timestamp and step
            all_keys = set()
            for entry in self.metrics_history:
                all_keys.update(entry.keys())
            metric_names = [key for key in all_keys if key not in ["timestamp", "step"]]
            
        if not metric_names:
            fig, ax = plt.subplots(figsize=figsize)
            ax.text(0.5, 0.5, "No metrics to plot", ha='center', va='center')
            return fig
            
        # Create the plot
        fig, ax = plt.subplots(figsize=figsize)
        
        # Get steps or use indices
        steps = []
        for entry in self.metrics_history:
            if "step" in entry and entry["step"] is not None:
                steps.append(entry["step"])
            else:
                steps.append(len(steps))
                
        # Plot each metric
        for metric_name in metric_names:
            values = [entry.get(metric_name, None) for entry in self.metrics_history]
            # Filter out None values
            valid_indices = [i for i, v in enumerate(values) if v is not None]
            valid_steps = [steps[i] for i in valid_indices]
            valid_values = [values[i] for i in valid_indices]
            
            if valid_values:
                ax.plot(valid_steps, valid_values, marker='o', linestyle='-', label=metric_name)
                
        ax.set_xlabel("Step")
        ax.set_ylabel("Value")
        ax.set_title("Metrics History")
        ax.legend()
        ax.grid(True)
        
        plt.tight_layout()
        return fig
    
    def plot_hyperparameter_importance(self, metric_name: str, figsize: Tuple[int, int] = (10, 6)) -> plt.Figure:
        """
        Plot hyperparameter importance for a specific metric.
        
        Args:
            metric_name: Name of the metric to analyze
            figsize: Figure size
            
        Returns:
            Matplotlib figure
        """
        try:
            from neural.hpo.parameter_importance import ParameterImportanceAnalyzer
            
            # Create a list of trials from metrics history and hyperparameters
            trials = []
            for i, metrics_entry in enumerate(self.metrics_history):
                if metric_name in metrics_entry:
                    trial = {
                        "parameters": self.hyperparameters,
                        "score": metrics_entry[metric_name]
                    }
                    trials.append(trial)
                    
            if not trials:
                fig, ax = plt.subplots(figsize=figsize)
                ax.text(0.5, 0.5, f"No data available for metric: {metric_name}", ha='center', va='center')
                return fig
                
            # Analyze parameter importance
            analyzer = ParameterImportanceAnalyzer()
            importance_dict = analyzer.analyze(trials, target_metric=metric_name)
            
            # Plot importance
            return analyzer.plot_importance(importance_dict, title=f"Hyperparameter Importance for {metric_name}", figsize=figsize)
            
        except (ImportError, Exception) as e:
            logger.error(f"Error plotting hyperparameter importance: {str(e)}")
            fig, ax = plt.subplots(figsize=figsize)
            ax.text(0.5, 0.5, f"Error plotting hyperparameter importance: {str(e)}", ha='center', va='center')
            return fig
    
    def save_experiment_summary(self) -> str:
        """
        Save a summary of the experiment.
        
        Returns:
            Path to the summary file
        """
        summary = {
            "experiment_name": self.experiment_name,
            "experiment_id": self.experiment_id,
            "metadata": self.metadata,
            "hyperparameters": self.hyperparameters,
            "metrics": {
                "latest": {k: v for entry in self.metrics_history[-1:] for k, v in entry.items() if k not in ["timestamp", "step"]} if self.metrics_history else {},
                "best": {}
            },
            "artifacts": list(self.artifacts.keys())
        }
        
        # Add best metrics
        all_metrics = set()
        for entry in self.metrics_history:
            all_metrics.update([k for k in entry.keys() if k not in ["timestamp", "step"]])
            
        for metric in all_metrics:
            best_value, step = self.get_best_metric(metric, mode="max")
            if best_value is not None:
                summary["metrics"]["best"][metric] = {
                    "value": best_value,
                    "step": step,
                    "mode": "max"
                }
                
        # Save the summary
        summary_path = os.path.join(self.experiment_dir, "summary.json")
        with open(summary_path, 'w') as f:
            json.dump(summary, f, indent=2)
            
        return summary_path
    
    def _save_metadata(self):
        """Save metadata to disk."""
        metadata_path = os.path.join(self.experiment_dir, "metadata.json")
        with open(metadata_path, 'w') as f:
            json.dump(self.metadata, f, indent=2)
    
    def _save_hyperparameters(self):
        """Save hyperparameters to disk."""
        hyperparameters_path = os.path.join(self.experiment_dir, "hyperparameters.json")
        with open(hyperparameters_path, 'w') as f:
            json.dump(self.hyperparameters, f, indent=2)
    
    def _save_metrics(self):
        """Save metrics history to disk."""
        metrics_path = os.path.join(self.experiment_dir, "metrics.json")
        with open(metrics_path, 'w') as f:
            json.dump(self.metrics_history, f, indent=2)
    
    def _save_artifacts(self):
        """Save artifacts metadata to disk."""
        artifacts_path = os.path.join(self.experiment_dir, "artifacts.json")
        with open(artifacts_path, 'w') as f:
            json.dump(self.artifacts, f, indent=2)
    
    def _get_artifact_type(self, artifact_path: str) -> str:
        """
        Determine the type of an artifact based on its file extension.
        
        Args:
            artifact_path: Path to the artifact
            
        Returns:
            Type of the artifact
        """
        ext = os.path.splitext(artifact_path)[1].lower()
        
        if ext in ['.h5', '.pt', '.pth', '.onnx', '.pb', '.tflite']:
            return 'model'
        elif ext in ['.png', '.jpg', '.jpeg', '.svg', '.gif']:
            return 'image'
        elif ext in ['.csv', '.tsv']:
            return 'table'
        elif ext in ['.json', '.yaml', '.yml']:
            return 'config'
        elif ext in ['.txt', '.log', '.md']:
            return 'text'
        else:
            return 'other'


class ExperimentManager:
    """Manages multiple experiments."""
    
    def __init__(self, base_dir: str = "neural_experiments"):
        """
        Initialize the experiment manager.
        
        Args:
            base_dir: Base directory for storing experiment data
        """
        self.base_dir = base_dir
        os.makedirs(base_dir, exist_ok=True)
        
    def create_experiment(self, experiment_name: str = None) -> ExperimentTracker:
        """
        Create a new experiment.
        
        Args:
            experiment_name: Name of the experiment (defaults to timestamp if None)
            
        Returns:
            ExperimentTracker instance
        """
        return ExperimentTracker(experiment_name=experiment_name, base_dir=self.base_dir)
    
    def get_experiment(self, experiment_id: str) -> Optional[ExperimentTracker]:
        """
        Get an existing experiment by ID.
        
        Args:
            experiment_id: ID of the experiment
            
        Returns:
            ExperimentTracker instance or None if not found
        """
        # Find the experiment directory
        for item in os.listdir(self.base_dir):
            if os.path.isdir(os.path.join(self.base_dir, item)) and item.endswith(f"_{experiment_id}"):
                # Extract experiment name
                experiment_name = item[:-(len(experiment_id) + 1)]
                
                # Create a new tracker and load the data
                tracker = ExperimentTracker(experiment_name=experiment_name, base_dir=self.base_dir)
                tracker.experiment_id = experiment_id
                tracker.experiment_dir = os.path.join(self.base_dir, item)
                
                # Load metadata
                metadata_path = os.path.join(tracker.experiment_dir, "metadata.json")
                if os.path.exists(metadata_path):
                    with open(metadata_path, 'r') as f:
                        tracker.metadata = json.load(f)
                
                # Load hyperparameters
                hyperparameters_path = os.path.join(tracker.experiment_dir, "hyperparameters.json")
                if os.path.exists(hyperparameters_path):
                    with open(hyperparameters_path, 'r') as f:
                        tracker.hyperparameters = json.load(f)
                
                # Load metrics
                metrics_path = os.path.join(tracker.experiment_dir, "metrics.json")
                if os.path.exists(metrics_path):
                    with open(metrics_path, 'r') as f:
                        tracker.metrics_history = json.load(f)
                
                # Load artifacts
                artifacts_path = os.path.join(tracker.experiment_dir, "artifacts.json")
                if os.path.exists(artifacts_path):
                    with open(artifacts_path, 'r') as f:
                        tracker.artifacts = json.load(f)
                
                return tracker
                
        return None
    
    def list_experiments(self) -> List[Dict[str, Any]]:
        """
        List all experiments.
        
        Returns:
            List of experiment summaries
        """
        experiments = []
        
        for item in os.listdir(self.base_dir):
            if os.path.isdir(os.path.join(self.base_dir, item)):
                # Check if it's an experiment directory
                metadata_path = os.path.join(self.base_dir, item, "metadata.json")
                if os.path.exists(metadata_path):
                    # Extract experiment ID
                    experiment_id = item.split('_')[-1]
                    
                    # Load metadata
                    with open(metadata_path, 'r') as f:
                        metadata = json.load(f)
                    
                    # Load summary if available
                    summary_path = os.path.join(self.base_dir, item, "summary.json")
                    summary = None
                    if os.path.exists(summary_path):
                        with open(summary_path, 'r') as f:
                            summary = json.load(f)
                    
                    # Create experiment summary
                    experiment = {
                        "experiment_name": item[:-(len(experiment_id) + 1)],
                        "experiment_id": experiment_id,
                        "status": metadata.get("status", "unknown"),
                        "start_time": metadata.get("start_time", "unknown"),
                        "end_time": metadata.get("end_time", None),
                        "summary": summary
                    }
                    
                    experiments.append(experiment)
        
        # Sort by start time (newest first)
        experiments.sort(key=lambda x: x.get("start_time", ""), reverse=True)
        
        return experiments
    
    def delete_experiment(self, experiment_id: str) -> bool:
        """
        Delete an experiment.
        
        Args:
            experiment_id: ID of the experiment
            
        Returns:
            True if successful, False otherwise
        """
        # Find the experiment directory
        for item in os.listdir(self.base_dir):
            if os.path.isdir(os.path.join(self.base_dir, item)) and item.endswith(f"_{experiment_id}"):
                # Delete the directory
                import shutil
                shutil.rmtree(os.path.join(self.base_dir, item))
                return True
                
        return False
    
    def compare_experiments(self, experiment_ids: List[str], metric_names: List[str] = None) -> Dict[str, plt.Figure]:
        """
        Compare multiple experiments.
        
        Args:
            experiment_ids: List of experiment IDs to compare
            metric_names: List of metric names to compare (compares all metrics if None)
            
        Returns:
            Dictionary mapping plot names to matplotlib figures
        """
        # Get experiments
        experiments = [self.get_experiment(exp_id) for exp_id in experiment_ids]
        experiments = [exp for exp in experiments if exp is not None]
        
        if not experiments:
            return {}
            
        # Determine which metrics to compare
        if metric_names is None:
            # Get all metric names from all experiments
            all_metrics = set()
            for exp in experiments:
                for entry in exp.metrics_history:
                    all_metrics.update([k for k in entry.keys() if k not in ["timestamp", "step"]])
            metric_names = sorted(list(all_metrics))
            
        if not metric_names:
            return {}
            
        # Create comparison plots
        plots = {}
        
        # Metrics comparison
        for metric_name in metric_names:
            fig, ax = plt.subplots(figsize=(10, 6))
            
            for exp in experiments:
                # Get steps and values for this metric
                steps = []
                values = []
                
                for entry in exp.metrics_history:
                    if "step" in entry and entry["step"] is not None and metric_name in entry:
                        steps.append(entry["step"])
                        values.append(entry[metric_name])
                        
                if steps and values:
                    ax.plot(steps, values, marker='o', linestyle='-', label=f"{exp.experiment_name} ({exp.experiment_id})")
                    
            ax.set_xlabel("Step")
            ax.set_ylabel(metric_name)
            ax.set_title(f"Comparison of {metric_name}")
            ax.legend()
            ax.grid(True)
            
            plt.tight_layout()
            plots[f"comparison_{metric_name}"] = fig
            
        # Hyperparameter comparison
        if len(experiments) > 1:
            # Get all hyperparameters from all experiments
            all_hyperparams = set()
            for exp in experiments:
                all_hyperparams.update(exp.hyperparameters.keys())
                
            if all_hyperparams:
                # Create a table of hyperparameters
                fig, ax = plt.subplots(figsize=(12, len(all_hyperparams) * 0.5 + 2))
                
                # Hide axes
                ax.axis('tight')
                ax.axis('off')
                
                # Create table data
                table_data = []
                header = ["Hyperparameter"] + [f"{exp.experiment_name} ({exp.experiment_id})" for exp in experiments]
                table_data.append(header)
                
                for param in sorted(all_hyperparams):
                    row = [param]
                    for exp in experiments:
                        value = exp.hyperparameters.get(param, "N/A")
                        row.append(str(value))
                    table_data.append(row)
                    
                # Create table
                table = ax.table(cellText=table_data[1:], colLabels=table_data[0], loc='center', cellLoc='center')
                table.auto_set_font_size(False)
                table.set_fontsize(10)
                table.scale(1, 1.5)
                
                ax.set_title("Hyperparameter Comparison")
                plt.tight_layout()
                
                plots["hyperparameter_comparison"] = fig
                
        return plots
