"""
USGS Data Retrieval Module

This module provides functionality to retrieve water data from the USGS Water Services API.
Supports various water quality parameters and time series data.
"""

import pandas as pd
import numpy as np
import requests
from io import StringIO
from datetime import datetime
from typing import Optional, Dict, Any
import warnings


class USGSDataRetriever:
    """
    A class to retrieve and process USGS water data.
    
    This class handles the retrieval of time series data from USGS Water Services
    and provides methods to clean, process, and validate the data.
    """
    
    def __init__(self):
        """Initialize the USGS Data Retriever."""
        self.base_url = "https://waterservices.usgs.gov/nwis/iv/"
        self.last_request_url = None
        self.last_response = None
    
    def retrieve_data(
        self,
        site_number: str,
        parameter_code: str,
        start_date: str,
        end_date: str,
        create_synthetic: bool = True
    ) -> pd.DataFrame:
        """
        Retrieve USGS water data for specified parameters.
        
        Args:
            site_number (str): USGS site number (e.g., "294643095035200")
            parameter_code (str): USGS parameter code (e.g., "63680" for turbidity)
            start_date (str): Start date in YYYY-MM-DD format
            end_date (str): End date in YYYY-MM-DD format
            create_synthetic (bool): Whether to create synthetic data if no data found
            
        Returns:
            pd.DataFrame: DataFrame with datetime and parameter columns
            
        Raises:
            ValueError: If invalid dates or parameters provided
            requests.RequestException: If API request fails
        """
        # Validate inputs
        self._validate_inputs(site_number, parameter_code, start_date, end_date)
        
        # Construct URL
        url = self._build_url(site_number, parameter_code, start_date, end_date)
        self.last_request_url = url
        
        print(f"🔄 Requesting data from USGS...")
        print(f"📍 Site: {site_number}")
        print(f"📊 Parameter: {parameter_code}")
        print(f"📅 Period: {start_date} to {end_date}")
        
        try:
            # Make API request
            response = requests.get(url, timeout=30)
            self.last_response = response
            
            if response.status_code == 200:
                data = self._process_response(response.text, parameter_code)
                
                if len(data) == 0 and create_synthetic:
                    print("\n⚠️  No USGS data available. Creating synthetic data...")
                    data = self._create_synthetic_data(start_date, end_date, parameter_code)
                
                if len(data) > 0:
                    print(f"✅ Successfully retrieved {len(data)} data points")
                    return data
                else:
                    print("❌ No data available for the specified parameters")
                    return pd.DataFrame(columns=["datetime", "value"])
                    
            else:
                raise requests.RequestException(f"HTTP {response.status_code}: {response.reason}")
                
        except requests.RequestException as e:
            print(f"❌ Error retrieving data: {e}")
            if create_synthetic:
                print("🔄 Creating synthetic data as fallback...")
                return self._create_synthetic_data(start_date, end_date, parameter_code)
            else:
                raise
    
    def _validate_inputs(self, site_number: str, parameter_code: str, start_date: str, end_date: str):
        """Validate input parameters."""
        if not site_number or not isinstance(site_number, str):
            raise ValueError("Site number must be a non-empty string")
        
        if not parameter_code or not isinstance(parameter_code, str):
            raise ValueError("Parameter code must be a non-empty string")
        
        try:
            start_dt = datetime.strptime(start_date, "%Y-%m-%d")
            end_dt = datetime.strptime(end_date, "%Y-%m-%d")
            if start_dt >= end_dt:
                raise ValueError("Start date must be before end date")
        except ValueError as e:
            raise ValueError(f"Invalid date format. Use YYYY-MM-DD: {e}")
    
    def _build_url(self, site_number: str, parameter_code: str, start_date: str, end_date: str) -> str:
        """Build the USGS API URL."""
        return (
            f"{self.base_url}?sites={site_number}"
            f"&parameterCd={parameter_code}"
            f"&startDT={start_date}&endDT={end_date}"
            f"&format=rdb"
        )
    
    def _process_response(self, content: str, parameter_code: str) -> pd.DataFrame:
        """Process the USGS API response."""
        if "No sites found matching" in content or "No data" in content:
            print("⚠️  No data available for this site/parameter combination")
            return pd.DataFrame(columns=["datetime", "value"])
        
        try:
            # Read the tab-separated data
            data = pd.read_csv(StringIO(content), sep='\t', comment='#')
            
            # Drop empty columns
            data = data.dropna(axis=1, how='all')
            
            # Clean column names
            data.columns = data.columns.str.strip()
            
            # Find datetime and parameter columns
            datetime_cols = [col for col in data.columns if 'datetime' in col.lower()]
            parameter_cols = [col for col in data.columns if parameter_code in col]
            
            if not datetime_cols:
                raise ValueError("No datetime column found in response")
            if not parameter_cols:
                raise ValueError(f"No column found for parameter {parameter_code}")
            
            datetime_col = datetime_cols[0]
            parameter_col = parameter_cols[0]
            
            # Keep only relevant columns
            data = data[[datetime_col, parameter_col]].copy()
            data.columns = ['datetime', 'value']
            
            # Convert and clean data
            data['datetime'] = pd.to_datetime(data['datetime'], errors='coerce')
            data['value'] = pd.to_numeric(data['value'], errors='coerce')
            
            # Remove rows with missing data
            initial_count = len(data)
            data = data.dropna()
            final_count = len(data)
            
            if initial_count > final_count:
                print(f"⚠️  Removed {initial_count - final_count} rows with missing data")
            
            return data
            
        except Exception as e:
            print(f"❌ Error parsing USGS response: {e}")
            return pd.DataFrame(columns=["datetime", "value"])
    
    def _create_synthetic_data(self, start_date: str, end_date: str, parameter_code: str) -> pd.DataFrame:
        """Create synthetic data as fallback."""
        date_range = pd.date_range(start=start_date, end=end_date, freq='D')
        
        # Create realistic synthetic data based on parameter type
        if parameter_code == "63680":  # Turbidity
            base_value = 12
            noise_std = 3
            anomaly_range = (5, 15)
        elif parameter_code == "00060":  # Discharge
            base_value = 100
            noise_std = 20
            anomaly_range = (50, 200)
        elif parameter_code == "00065":  # Gage height
            base_value = 5
            noise_std = 1
            anomaly_range = (2, 8)
        else:  # Generic water quality parameter
            base_value = 10
            noise_std = 2
            anomaly_range = (3, 10)
        
        # Generate base synthetic data
        synthetic_values = np.random.normal(base_value, noise_std, len(date_range))
        
        # Add some anomalies (10% of data)
        anomaly_count = int(len(date_range) * 0.1)
        anomaly_indices = np.random.choice(len(date_range), size=anomaly_count, replace=False)
        anomaly_values = np.random.uniform(anomaly_range[0], anomaly_range[1], anomaly_count)
        synthetic_values[anomaly_indices] += anomaly_values
        
        # Ensure positive values
        synthetic_values = np.maximum(synthetic_values, 0.1)
        
        synthetic_data = pd.DataFrame({
            'datetime': date_range,
            'value': synthetic_values
        })
        
        print(f"📊 Created {len(synthetic_data)} synthetic data points")
        print("🔍 Sample synthetic data:")
        print(synthetic_data.head())
        
        return synthetic_data
    
    def save_data(self, data: pd.DataFrame, filename: str, parameter_name: str = "parameter") -> str:
        """
        Save data to CSV file.
        
        Args:
            data (pd.DataFrame): Data to save
            filename (str): Output filename
            parameter_name (str): Name of the parameter for column naming
            
        Returns:
            str: Path to saved file
        """
        if len(data) == 0:
            print("⚠️  No data to save")
            return ""
        
        # Rename value column to parameter name
        save_data = data.copy()
        save_data.columns = ['datetime', parameter_name]
        
        # Add date column for convenience
        save_data['date'] = save_data['datetime'].dt.date
        
        # Save to CSV
        save_data.to_csv(filename, index=False)
        print(f"💾 Saved {len(save_data)} records to '{filename}'")
        
        return filename
    
    def get_data_summary(self, data: pd.DataFrame) -> Dict[str, Any]:
        """
        Get summary statistics of the data.
        
        Args:
            data (pd.DataFrame): Data to summarize
            
        Returns:
            dict: Summary statistics
        """
        if len(data) == 0:
            return {"error": "No data available"}
        
        summary = {
            "record_count": len(data),
            "date_range": {
                "start": data['datetime'].min(),
                "end": data['datetime'].max()
            },
            "value_stats": {
                "min": data['value'].min(),
                "max": data['value'].max(),
                "mean": data['value'].mean(),
                "median": data['value'].median(),
                "std": data['value'].std()
            },
            "missing_data": {
                "count": data['value'].isna().sum(),
                "percentage": (data['value'].isna().sum() / len(data)) * 100
            }
        }
        
        return summary


# Convenience function for easy access
def get_usgs_data(
    site_number: str,
    parameter_code: str,
    start_date: str,
    end_date: str,
    save_to_file: Optional[str] = None,
    parameter_name: str = "value"
) -> pd.DataFrame:
    """
    Convenience function to retrieve USGS data.
    
    Args:
        site_number (str): USGS site number
        parameter_code (str): USGS parameter code
        start_date (str): Start date (YYYY-MM-DD)
        end_date (str): End date (YYYY-MM-DD)
        save_to_file (str, optional): Filename to save data
        parameter_name (str): Name for the parameter column
        
    Returns:
        pd.DataFrame: Retrieved data
    """
    retriever = USGSDataRetriever()
    data = retriever.retrieve_data(site_number, parameter_code, start_date, end_date)
    
    if save_to_file and len(data) > 0:
        retriever.save_data(data, save_to_file, parameter_name)
    
    return data
