"""
Experiment Tracker

Comprehensive experiment tracking system with MLflow integration,
custom logging, and experiment management capabilities.
"""

import os
import json
import pickle
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Union, Callable
from datetime import datetime
import hashlib
import uuid
from pathlib import Path
import warnings

# Try to import MLflow
try:
    import mlflow
    import mlflow.sklearn
    import mlflow.pytorch
    import mlflow.tensorflow
    MLFLOW_AVAILABLE = True
except ImportError:
    MLFLOW_AVAILABLE = False

# Try to import other tracking libraries
try:
    import wandb
    WANDB_AVAILABLE = True
except ImportError:
    WANDB_AVAILABLE = False

try:
    import neptune
    NEPTUNE_AVAILABLE = True
except ImportError:
    NEPTUNE_AVAILABLE = False


class ExperimentTracker:
    """
    Comprehensive experiment tracking system supporting multiple backends
    including MLflow, Weights & Biases, and Neptune.
    """
    
    def __init__(self, tracking_uri: Optional[str] = None,
                 experiment_name: Optional[str] = None,
                 backend: str = "mlflow"):
        """
        Initialize experiment tracker
        
        Args:
            tracking_uri: URI for tracking server (MLflow)
            experiment_name: Name of the experiment
            backend: Tracking backend ('mlflow', 'wandb', 'neptune', 'local')
        """
        self.backend = backend.lower()
        self.experiment_name = experiment_name or f"experiment_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.tracking_uri = tracking_uri
        self.current_run = None
        self.run_history = []
        self.artifacts = []
        self.metrics_history = {}
        
        # Initialize backend
        self._initialize_backend()
        
        # Local storage fallback
        self.local_storage = Path("./experiments")
        self.local_storage.mkdir(exist_ok=True)
        self.experiment_dir = self.local_storage / self.experiment_name
        self.experiment_dir.mkdir(exist_ok=True)
    
    def _initialize_backend(self):
        """Initialize the tracking backend"""
        if self.backend == "mlflow":
            if not MLFLOW_AVAILABLE:
                print("MLflow not available. Install with: pip install mlflow")
                self.backend = "local"
                return
            
            if self.tracking_uri:
                mlflow.set_tracking_uri(self.tracking_uri)
            
            mlflow.set_experiment(self.experiment_name)
            print(f"MLflow experiment initialized: {self.experiment_name}")
        
        elif self.backend == "wandb":
            if not WANDB_AVAILABLE:
                print("Weights & Biases not available. Install with: pip install wandb")
                self.backend = "local"
                return
            
            wandb.init(project=self.experiment_name)
            print(f"W&B experiment initialized: {self.experiment_name}")
        
        elif self.backend == "neptune":
            if not NEPTUNE_AVAILABLE:
                print("Neptune not available. Install with: pip install neptune")
                self.backend = "local"
                return
            
            # Neptune initialization would require API token
            print("Neptune backend requires API token configuration")
            self.backend = "local"
        
        else:
            print(f"Using local storage backend: {self.experiment_dir}")
    
    def start_run(self, run_name: Optional[str] = None, 
                 tags: Optional[Dict[str, str]] = None) -> str:
        """
        Start a new experiment run
        
        Args:
            run_name: Name for the run
            tags: Tags for the run
            
        Returns:
            Run ID
        """
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            self.current_run = mlflow.start_run(run_name=run_name, tags=tags)
            run_id = self.current_run.info.run_id
            print(f"MLflow run started: {run_id}")
        
        elif self.backend == "wandb" and WANDB_AVAILABLE:
            run_id = wandb.run.id
            if run_name:
                wandb.run.name = run_name
            if tags:
                wandb.config.update(tags)
            print(f"W&B run active: {run_id}")
        
        else:
            # Local storage
            run_id = str(uuid.uuid4())[:8]
            run_dir = self.experiment_dir / run_id
            run_dir.mkdir(exist_ok=True)
            
            self.current_run = {
                'run_id': run_id,
                'run_name': run_name or f"run_{run_id}",
                'start_time': datetime.now().isoformat(),
                'tags': tags or {},
                'run_dir': run_dir,
                'metrics': {},
                'params': {},
                'artifacts': []
            }
            
            # Save run metadata
            with open(run_dir / "run_metadata.json", 'w') as f:
                json.dump(self.current_run, f, indent=2)
            
            print(f"Local run started: {run_id}")
        
        self.run_history.append({
            'run_id': run_id,
            'run_name': run_name,
            'start_time': datetime.now().isoformat(),
            'backend': self.backend
        })
        
        return run_id
    
    def log_params(self, params: Dict[str, Any]):
        """Log parameters to the current run"""
        if not self.current_run:
            raise ValueError("No active run. Start a run first.")
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            for key, value in params.items():
                mlflow.log_param(key, value)
        
        elif self.backend == "wandb" and WANDB_AVAILABLE:
            wandb.config.update(params)
        
        else:
            # Local storage
            self.current_run['params'].update(params)
            self._save_local_run()
        
        print(f"Logged {len(params)} parameters")
    
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
        """Log metrics to the current run"""
        if not self.current_run:
            raise ValueError("No active run. Start a run first.")
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            for key, value in metrics.items():
                mlflow.log_metric(key, value, step=step)
        
        elif self.backend == "wandb" and WANDB_AVAILABLE:
            wandb.log(metrics, step=step)
        
        else:
            # Local storage
            if step is not None:
                for key, value in metrics.items():
                    if key not in self.current_run['metrics']:
                        self.current_run['metrics'][key] = []
                    self.current_run['metrics'][key].append({
                        'step': step,
                        'value': value,
                        'timestamp': datetime.now().isoformat()
                    })
            else:
                self.current_run['metrics'].update(metrics)
            
            self._save_local_run()
        
        # Update local metrics history
        for key, value in metrics.items():
            if key not in self.metrics_history:
                self.metrics_history[key] = []
            self.metrics_history[key].append({
                'value': value,
                'step': step,
                'timestamp': datetime.now().isoformat()
            })
        
        print(f"Logged {len(metrics)} metrics")
    
    def log_artifact(self, artifact_path: str, artifact_name: Optional[str] = None):
        """Log an artifact to the current run"""
        if not self.current_run:
            raise ValueError("No active run. Start a run first.")
        
        if not os.path.exists(artifact_path):
            raise FileNotFoundError(f"Artifact not found: {artifact_path}")
        
        name = artifact_name or os.path.basename(artifact_path)
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            mlflow.log_artifact(artifact_path, artifact_path=name)
        
        elif self.backend == "wandb" and WANDB_AVAILABLE:
            wandb.save(artifact_path, base_path=wandb.run.dir)
        
        else:
            # Local storage
            import shutil
            run_dir = self.current_run['run_dir']
            artifact_dest = run_dir / "artifacts" / name
            artifact_dest.parent.mkdir(exist_ok=True)
            shutil.copy2(artifact_path, artifact_dest)
            
            self.current_run['artifacts'].append({
                'name': name,
                'original_path': artifact_path,
                'stored_path': str(artifact_dest),
                'timestamp': datetime.now().isoformat()
            })
            
            self._save_local_run()
        
        self.artifacts.append({
            'name': name,
            'path': artifact_path,
            'timestamp': datetime.now().isoformat()
        })
        
        print(f"Logged artifact: {name}")
    
    def log_model(self, model, model_name: str, framework: str = "sklearn"):
        """Log a model to the current run"""
        if not self.current_run:
            raise ValueError("No active run. Start a run first.")
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            if framework == "sklearn":
                mlflow.sklearn.log_model(model, model_name)
            elif framework == "pytorch":
                mlflow.pytorch.log_model(model, model_name)
            elif framework == "tensorflow":
                mlflow.tensorflow.log_model(model, model_name)
            else:
                # Generic model saving
                model_path = f"{model_name}.pkl"
                with open(model_path, 'wb') as f:
                    pickle.dump(model, f)
                mlflow.log_artifact(model_path, f"models/{model_name}")
                os.remove(model_path)
        
        else:
            # Local storage
            run_dir = self.current_run['run_dir']
            models_dir = run_dir / "models"
            models_dir.mkdir(exist_ok=True)
            
            model_path = models_dir / f"{model_name}.pkl"
            with open(model_path, 'wb') as f:
                pickle.dump(model, f)
            
            self.current_run['artifacts'].append({
                'name': model_name,
                'type': 'model',
                'framework': framework,
                'path': str(model_path),
                'timestamp': datetime.now().isoformat()
            })
            
            self._save_local_run()
        
        print(f"Logged model: {model_name} ({framework})")
    
    def log_dataset(self, dataset: pd.DataFrame, dataset_name: str):
        """Log a dataset to the current run"""
        if not self.current_run:
            raise ValueError("No active run. Start a run first.")
        
        # Save dataset info
        dataset_info = {
            'name': dataset_name,
            'shape': dataset.shape,
            'columns': list(dataset.columns),
            'dtypes': dataset.dtypes.to_dict(),
            'missing_values': dataset.isnull().sum().to_dict(),
            'timestamp': datetime.now().isoformat()
        }
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            # Log dataset info as artifact
            info_path = f"{dataset_name}_info.json"
            with open(info_path, 'w') as f:
                json.dump(dataset_info, f, indent=2)
            mlflow.log_artifact(info_path, f"datasets/{dataset_name}")
            os.remove(info_path)
            
            # Log dataset as artifact
            dataset_path = f"{dataset_name}.csv"
            dataset.to_csv(dataset_path, index=False)
            mlflow.log_artifact(dataset_path, f"datasets/{dataset_name}")
            os.remove(dataset_path)
        
        else:
            # Local storage
            run_dir = self.current_run['run_dir']
            datasets_dir = run_dir / "datasets"
            datasets_dir.mkdir(exist_ok=True)
            
            # Save dataset info
            info_path = datasets_dir / f"{dataset_name}_info.json"
            with open(info_path, 'w') as f:
                json.dump(dataset_info, f, indent=2)
            
            # Save dataset
            dataset_path = datasets_dir / f"{dataset_name}.csv"
            dataset.to_csv(dataset_path, index=False)
            
            self.current_run['artifacts'].append({
                'name': dataset_name,
                'type': 'dataset',
                'info_path': str(info_path),
                'data_path': str(dataset_path),
                'timestamp': datetime.now().isoformat()
            })
            
            self._save_local_run()
        
        print(f"Logged dataset: {dataset_name}")
    
    def end_run(self, status: str = "FINISHED"):
        """End the current run"""
        if not self.current_run:
            print("No active run to end")
            return
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            mlflow.end_run()
            print(f"MLflow run ended: {self.current_run.info.run_id}")
        
        elif self.backend == "wandb" and WANDB_AVAILABLE:
            wandb.finish()
            print(f"W&B run ended: {wandb.run.id}")
        
        else:
            # Local storage
            self.current_run['end_time'] = datetime.now().isoformat()
            self.current_run['status'] = status
            self._save_local_run()
            print(f"Local run ended: {self.current_run['run_id']}")
        
        self.current_run = None
    
    def _save_local_run(self):
        """Save current run to local storage"""
        if self.backend != "local" or not self.current_run:
            return
        
        run_dir = self.current_run['run_dir']
        with open(run_dir / "run_metadata.json", 'w') as f:
            json.dump(self.current_run, f, indent=2)
    
    def list_runs(self, experiment_name: Optional[str] = None) -> pd.DataFrame:
        """List all runs in the experiment"""
        exp_name = experiment_name or self.experiment_name
        
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            from mlflow.tracking import MlflowClient
            client = MlflowClient()
            runs = client.search_runs(experiment_ids=[client.get_experiment_by_name(exp_name).experiment_id])
            
            run_data = []
            for run in runs:
                run_data.append({
                    'run_id': run.info.run_id,
                    'name': run.info.run_name,
                    'status': run.info.status,
                    'start_time': run.info.start_time,
                    'end_time': run.info.end_time,
                    'metrics': run.data.metrics,
                    'params': run.data.params
                })
            
            return pd.DataFrame(run_data)
        
        else:
            # Local storage
            exp_dir = self.local_storage / exp_name
            if not exp_dir.exists():
                return pd.DataFrame()
            
            run_data = []
            for run_dir in exp_dir.iterdir():
                if run_dir.is_dir():
                    metadata_path = run_dir / "run_metadata.json"
                    if metadata_path.exists():
                        with open(metadata_path, 'r') as f:
                            metadata = json.load(f)
                        
                        run_data.append({
                            'run_id': metadata.get('run_id'),
                            'name': metadata.get('run_name'),
                            'status': metadata.get('status', 'ACTIVE'),
                            'start_time': metadata.get('start_time'),
                            'end_time': metadata.get('end_time'),
                            'metrics': metadata.get('metrics', {}),
                            'params': metadata.get('params', {})
                        })
            
            return pd.DataFrame(run_data)
    
    def get_run(self, run_id: str) -> Dict[str, Any]:
        """Get details of a specific run"""
        if self.backend == "mlflow" and MLFLOW_AVAILABLE:
            from mlflow.tracking import MlflowClient
            client = MlflowClient()
            run = client.get_run(run_id)
            
            return {
                'run_id': run.info.run_id,
                'name': run.info.run_name,
                'status': run.info.status,
                'start_time': run.info.start_time,
                'end_time': run.info.end_time,
                'metrics': run.data.metrics,
                'params': run.data.params,
                'tags': run.data.tags
            }
        
        else:
            # Local storage
            run_dir = self.experiment_dir / run_id
            metadata_path = run_dir / "run_metadata.json"
            
            if metadata_path.exists():
                with open(metadata_path, 'r') as f:
                    return json.load(f)
            else:
                raise FileNotFoundError(f"Run {run_id} not found")
    
    def compare_runs(self, run_ids: List[str], metric: str) -> pd.DataFrame:
        """Compare runs based on a specific metric"""
        comparison_data = []
        
        for run_id in run_ids:
            run_info = self.get_run(run_id)
            metrics = run_info.get('metrics', {})
            
            comparison_data.append({
                'run_id': run_id,
                'name': run_info.get('name', run_id),
                'metric_value': metrics.get(metric, None),
                'status': run_info.get('status', 'UNKNOWN')
            })
        
        df = pd.DataFrame(comparison_data)
        
        if metric in df.columns and df['metric_value'].notna().any():
            df = df.sort_values('metric_value', ascending=False)
        
        return df
    
    def plot_metrics_comparison(self, run_ids: List[str], metrics: List[str]):
        """Plot comparison of metrics across runs"""
        import matplotlib.pyplot as plt
        
        fig, axes = plt.subplots(len(metrics), 1, figsize=(10, 5 * len(metrics)))
        if len(metrics) == 1:
            axes = [axes]
        
        for i, metric in enumerate(metrics):
            metric_values = []
            run_names = []
            
            for run_id in run_ids:
                run_info = self.get_run(run_id)
                run_metrics = run_info.get('metrics', {})
                
                if metric in run_metrics:
                    metric_values.append(run_metrics[metric])
                    run_names.append(run_info.get('name', run_id))
            
            if metric_values:
                axes[i].bar(run_names, metric_values)
                axes[i].set_title(f'{metric} Comparison')
                axes[i].set_ylabel(metric)
                axes[i].tick_params(axis='x', rotation=45)
        
        plt.tight_layout()
        plt.show()
    
    def get_best_run(self, metric: str, experiment_name: Optional[str] = None) -> Dict[str, Any]:
        """Get the best run based on a specific metric"""
        runs_df = self.list_runs(experiment_name)
        
        if runs_df.empty:
            return {}
        
        # Filter runs that have the metric
        valid_runs = runs_df[runs_df['metrics'].apply(lambda x: metric in x if isinstance(x, dict) else False)]
        
        if valid_runs.empty:
            return {}
        
        # Find best run
        best_run = None
        best_value = None
        
        for _, run in valid_runs.iterrows():
            value = run['metrics'][metric]
            if best_value is None or value > best_value:
                best_value = value
                best_run = run.to_dict()
        
        return best_run
    
    def export_experiment(self, output_path: str, experiment_name: Optional[str] = None):
        """Export experiment data to file"""
        exp_name = experiment_name or self.experiment_name
        
        export_data = {
            'experiment_name': exp_name,
            'export_time': datetime.now().isoformat(),
            'backend': self.backend,
            'runs': []
        }
        
        runs_df = self.list_runs(exp_name)
        
        for _, run in runs_df.iterrows():
            run_details = self.get_run(run['run_id'])
            export_data['runs'].append(run_details)
        
        with open(output_path, 'w') as f:
            json.dump(export_data, f, indent=2)
        
        print(f"Experiment exported to: {output_path}")
    
    def import_experiment(self, import_path: str):
        """Import experiment data from file"""
        with open(import_path, 'r') as f:
            import_data = json.load(f)
        
        # This would need to be adapted based on the backend
        print(f"Importing experiment: {import_data['experiment_name']}")
        print(f"Runs to import: {len(import_data['runs'])}")
        
        # Implementation would depend on backend capabilities
        # For now, just show what would be imported
        for run in import_data['runs']:
            print(f"  - {run.get('name', run.get('run_id', 'Unknown'))}")
    
    def cleanup_old_runs(self, days_old: int = 30, keep_best: bool = True):
        """Clean up old runs from local storage"""
        if self.backend != "local":
            print("Cleanup only available for local storage backend")
            return
        
        cutoff_date = datetime.now() - pd.Timedelta(days=days_old)
        cleaned_runs = []
        
        for run_dir in self.experiment_dir.iterdir():
            if run_dir.is_dir():
                metadata_path = run_dir / "run_metadata.json"
                if metadata_path.exists():
                    with open(metadata_path, 'r') as f:
                        metadata = json.load(f)
                    
                    start_time = datetime.fromisoformat(metadata.get('start_time', ''))
                    
                    if start_time < cutoff_date:
                        # Check if we should keep this run
                        if keep_best:
                            # This is simplified - would need to check performance
                            pass
                        
                        # Remove run directory
                        import shutil
                        shutil.rmtree(run_dir)
                        cleaned_runs.append(metadata.get('run_id', run_dir.name))
        
        print(f"Cleaned up {len(cleaned_runs)} old runs")
        return cleaned_runs
    
    def get_experiment_summary(self, experiment_name: Optional[str] = None) -> Dict[str, Any]:
        """Get summary of the experiment"""
        exp_name = experiment_name or self.experiment_name
        runs_df = self.list_runs(exp_name)
        
        if runs_df.empty:
            return {"error": "No runs found"}
        
        summary = {
            'experiment_name': exp_name,
            'total_runs': len(runs_df),
            'backend': self.backend,
            'status_counts': runs_df['status'].value_counts().to_dict(),
            'date_range': {
                'earliest': runs_df['start_time'].min(),
                'latest': runs_df['start_time'].max()
            }
        }
        
        # Extract all metrics
        all_metrics = set()
        for metrics in runs_df['metrics']:
            if isinstance(metrics, dict):
                all_metrics.update(metrics.keys())
        
        summary['available_metrics'] = list(all_metrics)
        
        # Best runs for each metric
        summary['best_runs'] = {}
        for metric in all_metrics:
            best_run = self.get_best_run(metric, exp_name)
            if best_run:
                summary['best_runs'][metric] = {
                    'run_id': best_run['run_id'],
                    'name': best_run.get('name', best_run['run_id']),
                    'value': best_run['metrics'][metric]
                }
        
        return summary
