"""
GHCN data workflow for parameter estimation.

This module provides a complete workflow for:
- Loading historical precipitation data from GHCN
- Calculating simulation parameters
- Creating ParameterManifest objects for simulation

This is separate from the PrecipitationSimulator to maintain clean separation
between data gathering and simulation workflows.
"""

import pandas as pd
import numpy as np
from datetime import datetime
from typing import Dict, List, Optional, Union
import logging

from ..data.ghcn_downloader import GHCNDownloader, find_nearby_stations
from ..data.ghcn_parser import GHCNParser
from ..engines.analytical import AnalyticalEngine, ParameterManifest


class GHCNWorkflow:
    """
    Complete workflow for GHCN data processing and parameter estimation.
    
    This class handles the data gathering phase that produces ParameterManifest
    objects for use with PrecipitationSimulator.
    
    Example:
        >>> workflow = pg.GHCNWorkflow()
        >>> workflow.load_historical_data(latitude=40.7608, longitude=-111.8910)
        >>> parameters = workflow.calculate_parameters()
        >>> 
        >>> # Now use parameters with PrecipitationSimulator
        >>> sim = pg.PrecipitationSimulator(parameters=parameters)
        >>> results = sim.run_monte_carlo_simulation(n_realizations=100)
    """
    
    def __init__(self, station_id: Optional[str] = None, cache_dir: str = 'data'):
        """
        Initialize the GHCN workflow.
        
        Args:
            station_id: Optional GHCN station ID. If None, will find nearest station.
            cache_dir: Directory for caching downloaded data
        """
        self.station_id = station_id
        self.cache_dir = cache_dir
        self.historical_data = None
        self.station_info = {}
        
        # Set up logging
        self.logger = logging.getLogger(__name__)
        
    def load_historical_data(self, latitude: Optional[float] = None, 
                           longitude: Optional[float] = None,
                           radius_km: float = 50, min_years: int = 20) -> pd.Series:
        """
        Load and process historical precipitation data from GHCN.
        
        Args:
            latitude: Latitude for station search (if station_id not provided)
            longitude: Longitude for station search (if station_id not provided)
            radius_km: Search radius in kilometers
            min_years: Minimum years of data required
            
        Returns:
            Historical precipitation time series
            
        Raises:
            ValueError: If no suitable stations found or data cannot be loaded
        """
        if self.station_id is None:
            if latitude is None or longitude is None:
                raise ValueError("Must provide either station_id or latitude/longitude")
            
            # Find nearest station
            stations = find_nearby_stations(
                latitude=latitude, longitude=longitude, 
                radius_km=radius_km, min_years=min_years
            )
            
            if not stations:
                raise ValueError(f"No suitable stations found within {radius_km} km")
            
            stations_df = pd.DataFrame(stations)
            self.station_id = stations_df.iloc[0]['id']
            self.station_info = stations_df.iloc[0].to_dict()
            
            self.logger.info(f"Using station: {self.station_id} - {self.station_info['name']}")
            self.logger.info(f"Distance: {self.station_info['distance_km']:.1f} km")
            self.logger.info(f"Data period: {self.station_info['first_year']}-{self.station_info['last_year']}")
        
        # Download and parse data
        downloader = GHCNDownloader(cache_dir=self.cache_dir)
        dly_path = downloader.download_station_data(self.station_id)
        
        parser = GHCNParser(dly_path)
        ghcn_data = parser.parse_dly_file(dly_path)
        self.historical_data = parser.extract_precipitation(ghcn_data).sort_index().astype(float)
        
        self.logger.info(f"Historical data loaded: {len(self.historical_data)} days")
        self.logger.info(f"Period: {self.historical_data.index.min().date()} to {self.historical_data.index.max().date()}")
        
        return self.historical_data
    
    def calculate_parameters(self, 
                           threshold: float = 0.1,
                           enable_trend_analysis: bool = False,
                           sliding_window_years: int = 20) -> ParameterManifest:
        """
        Calculate WGEN parameters from historical data using AnalyticalEngine.
        
        Args:
            threshold: Precipitation threshold for wet/dry classification (mm)
            enable_trend_analysis: Enable trend detection and analysis
            sliding_window_years: Window size for sliding window analysis
            
        Returns:
            ParameterManifest ready for simulation
            
        Raises:
            ValueError: If historical data not loaded
        """
        if self.historical_data is None:
            raise ValueError("Historical data not loaded. Call load_historical_data() first.")
        
        # Create analytical engine
        engine = AnalyticalEngine()
        
        # Load data into engine
        engine.load_data(self.historical_data)
        
        # Calculate parameters
        parameters = engine.calculate_monthly_parameters(threshold=threshold)
        
        # Perform trend analysis if requested
        if enable_trend_analysis:
            self.logger.info("Performing trend analysis...")
            trend_analysis = engine.detect_trends()
            parameters.trend_analysis = trend_analysis
        
        # Add metadata
        parameters.metadata.update({
            'station_id': self.station_id,
            'station_info': self.station_info,
            'data_period': {
                'start': self.historical_data.index.min().isoformat(),
                'end': self.historical_data.index.max().isoformat(),
                'total_days': len(self.historical_data)
            },
            'threshold': threshold,
            'trend_analysis_enabled': enable_trend_analysis,
            'created_at': datetime.now().isoformat()
        })
        
        self.logger.info("Parameter calculation completed successfully")
        return parameters
    
    def get_data_summary(self) -> Dict:
        """
        Get summary statistics of the loaded historical data.
        
        Returns:
            Dictionary with data summary statistics
        """
        if self.historical_data is None:
            raise ValueError("Historical data not loaded")
        
        data = self.historical_data.dropna()
        
        return {
            'total_days': len(data),
            'date_range': {
                'start': data.index.min().date(),
                'end': data.index.max().date()
            },
            'annual_statistics': {
                'mean': data.resample('YE').sum().mean(),
                'std': data.resample('YE').sum().std(),
                'min': data.resample('YE').sum().min(),
                'max': data.resample('YE').sum().max()
            },
            'wet_day_statistics': {
                'frequency': (data > 0.1).mean(),
                'mean_wet_amount': data[data > 0.1].mean(),
                'max_daily': data.max()
            },
            'monthly_totals': data.groupby(data.index.month).sum().to_dict(),
            'station_info': self.station_info
        }
    
    def export_data(self, output_path: str) -> None:
        """
        Export historical data to CSV file.
        
        Args:
            output_path: Path to save the data
        """
        if self.historical_data is None:
            raise ValueError("Historical data not loaded")
        
        self.historical_data.to_csv(output_path)
        self.logger.info(f"Historical data exported to {output_path}")
    
    def create_parameter_manifest_from_data(self, 
                                          data: pd.Series,
                                          threshold: float = 0.1,
                                          enable_trend_analysis: bool = False) -> ParameterManifest:
        """
        Create a ParameterManifest from provided data (alternative to load_historical_data).
        
        Args:
            data: Historical precipitation time series
            threshold: Precipitation threshold for wet/dry classification (mm)
            enable_trend_analysis: Enable trend detection and analysis
            
        Returns:
            ParameterManifest ready for simulation
        """
        # Store the data
        self.historical_data = data
        
        # Calculate parameters using the standard workflow
        return self.calculate_parameters(
            threshold=threshold,
            enable_trend_analysis=enable_trend_analysis
        )