"""
SOCAT Data Loader
Surface Ocean CO₂ Atlas - Ocean pCO₂ measurements
"""

import csv
import os
from typing import Dict, List, Optional


class SOCATLoader:
    """
    Loader for SOCAT ocean CO₂ data
    """
    
    def __init__(self, data_dir: str = "./data"):
        """
        Initialize SOCAT loader
        
        Parameters
        ----------
        data_dir : str
            Base data directory
        """
        self.data_dir = data_dir
        self.socat_dir = f"{data_dir}/socat"
        self.data = []
        self._load_reference_data()
    
    def _load_reference_data(self):
        """Load reference data from research paper"""
        self.reference_data = [
            {'year': 1970, 'pco2_sw_mean': 320, 'sink_strength': -1.21},
            {'year': 1980, 'pco2_sw_mean': 330, 'sink_strength': -1.78},
            {'year': 1990, 'pco2_sw_mean': 340, 'sink_strength': -1.93},
            {'year': 2000, 'pco2_sw_mean': 355, 'sink_strength': -2.24},
            {'year': 2010, 'pco2_sw_mean': 370, 'sink_strength': -2.71},
            {'year': 2020, 'pco2_sw_mean': 390, 'sink_strength': -3.05},
            {'year': 2025, 'pco2_sw_mean': 400, 'sink_strength': -3.08}
        ]
        
        # Ocean basin data
        self.basin_data = {
            'north_atlantic': {
                'area': 41.5e6,  # km²
                'pco2_2025': 395,
                'sink_2025': -0.42
            },
            'south_atlantic': {
                'area': 40.3e6,
                'pco2_2025': 398,
                'sink_2025': -0.38
            },
            'north_pacific': {
                'area': 77.0e6,
                'pco2_2025': 405,
                'sink_2025': -0.65
            },
            'south_pacific': {
                'area': 84.5e6,
                'pco2_2025': 402,
                'sink_2025': -0.58
            },
            'indian': {
                'area': 70.6e6,
                'pco2_2025': 398,
                'sink_2025': -0.52
            },
            'southern': {
                'area': 21.9e6,
                'pco2_2025': 390,
                'sink_2025': -0.53
            }
        }
    
    def load_from_file(self, filepath: str) -> List[Dict]:
        """
        Load SOCAT data from CSV file
        
        Parameters
        ----------
        filepath : str
            Path to CSV file
        
        Returns
        -------
        list
            List of data records
        """
        if not os.path.exists(filepath):
            return self.reference_data
        
        data = []
        with open(filepath, 'r') as f:
            reader = csv.DictReader(f)
            for row in reader:
                record = {}
                for key, value in row.items():
                    try:
                        record[key] = float(value)
                    except ValueError:
                        record[key] = value
                data.append(record)
        
        self.data = data
        return data
    
    def get_global_sink(self, year: int) -> float:
        """
        Get global ocean sink strength for given year
        
        Parameters
        ----------
        year : int
            Target year
        
        Returns
        -------
        float
            Ocean sink strength (PgC/yr)
        """
        closest = min(self.reference_data, key=lambda x: abs(x['year'] - year))
        return closest['sink_strength']
    
    def get_basin_sink(self, basin: str, year: int = 2025) -> float:
        """
        Get sink strength for specific ocean basin
        
        Parameters
        ----------
        basin : str
            Ocean basin name
        year : int
            Target year
        
        Returns
        -------
        float
            Basin sink strength (PgC/yr)
        """
        if basin in self.basin_data:
            return self.basin_data[basin]['sink_2025']
        return 0.0
    
    def get_pco2_sw(self, year: int, basin: Optional[str] = None) -> float:
        """
        Get surface ocean pCO₂
        
        Parameters
        ----------
        year : int
            Target year
        basin : str, optional
            Ocean basin
        
        Returns
        -------
        float
            pCO₂ in µatm
        """
        if basin and basin in self.basin_data:
            return self.basin_data[basin]['pco2_2025']
        
        closest = min(self.reference_data, key=lambda x: abs(x['year'] - year))
        return closest['pco2_sw_mean']
    
    def get_delta_pco2(self, year: int, pco2_atm: Optional[float] = None) -> float:
        """
        Get air-sea pCO₂ difference
        
        Parameters
        ----------
        year : int
            Target year
        pco2_atm : float, optional
            Atmospheric pCO₂ (if None, uses Keeling value)
        
        Returns
        -------
        float
            ΔpCO₂ (atm - ocean) in µatm
        """
        pco2_sw = self.get_pco2_sw(year)
        
        if pco2_atm is None:
            # Approximate atmospheric pCO₂
            if year <= 2025:
                pco2_atm = 316.9 + (year - 1960) * 1.8
            else:
                pco2_atm = 424.0 + (year - 2025) * 2.4
        
        return pco2_atm - pco2_sw
    
    def summary(self) -> str:
        """Print summary of SOCAT data"""
        latest = self.reference_data[-1]
        
        summary = f"""
╔════════════════════════════════════════════════════════════════╗
║                      SOCAT Summary                             ║
╠════════════════════════════════════════════════════════════════╣
║  Latest year     : {latest['year']}                              ║
║  Global sink     : {latest['sink_strength']:.2f} PgC/yr                    ║
║  Mean surface pCO₂ : {latest['pco2_sw_mean']} µatm                         ║
╠════════════════════════════════════════════════════════════════╣
║  Basin sinks (2025):                                          ║
║    North Atlantic : {self.basin_data['north_atlantic']['sink_2025']:.2f} PgC/yr        ║
║    North Pacific  : {self.basin_data['north_pacific']['sink_2025']:.2f} PgC/yr        ║
║    Indian Ocean   : {self.basin_data['indian']['sink_2025']:.2f} PgC/yr        ║
║    Southern Ocean : {self.basin_data['southern']['sink_2025']:.2f} PgC/yr        ║
╚════════════════════════════════════════════════════════════════╝
        """
        return summary
