#!/usr/bin/env python3
"""
Comprehensive Arctic GRIB data extractor that bypasses ECCODES polar stereographic errors
Extracts lat/lon coordinates and data from every point in Arctic GRIB files
"""

import subprocess
import os
import tempfile
import numpy as np
import pandas as pd
import json
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

class ArcticGRIBExtractor:
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp()
        print(f"Using temporary directory: {self.temp_dir}")
    
    def check_tools(self):
        """Check if required tools are available"""
        tools = {
            'wgrib2': self._check_wgrib2(),
            'pygrib': self._check_pygrib(),
            'eccodes': self._check_eccodes()
        }
        
        print("Tool availability:")
        for tool, available in tools.items():
            status = "✓" if available else "✗"
            print(f"  {status} {tool}")
        
        return tools
    
    def _check_wgrib2(self):
        try:
            subprocess.run(['wgrib2', '-version'], capture_output=True, check=True)
            return True
        except:
            return False
    
    def _check_pygrib(self):
        try:
            import pygrib
            return True
        except ImportError:
            return False
    
    def _check_eccodes(self):
        try:
            import eccodes
            return True
        except ImportError:
            return False
    
    def extract_with_wgrib2_comprehensive(self, grib_file):
        """
        Extract all data using wgrib2 - most reliable method for polar stereographic
        """
        print("\n=== Method 1: wgrib2 comprehensive extraction ===")
        
        try:
            # First, get inventory of the GRIB file
            print("Getting GRIB inventory...")
            cmd_inv = ['wgrib2', grib_file, '-inv', '/dev/stdout']
            result = subprocess.run(cmd_inv, capture_output=True, text=True, check=True)
            inventory = result.stdout
            print(f"GRIB inventory:\n{inventory}")
            
            # Get grid information
            print("Getting grid information...")
            cmd_grid = ['wgrib2', grib_file, '-grid']
            result = subprocess.run(cmd_grid, capture_output=True, text=True, check=True)
            grid_info = result.stdout
            print(f"Grid info: {grid_info.strip()}")
            
            # Extract lat/lon and data using -csv option (most reliable)
            csv_file = os.path.join(self.temp_dir, 'arctic_data.csv')
            print(f"Extracting to CSV: {csv_file}")
            
            cmd_csv = ['wgrib2', grib_file, '-csv', csv_file]
            result = subprocess.run(cmd_csv, capture_output=True, text=True, check=True)
            
            if os.path.exists(csv_file) and os.path.getsize(csv_file) > 0:
                print("Successfully extracted to CSV!")
                
                # Read the CSV data
                df = pd.read_csv(csv_file)
                print(f"Loaded DataFrame with {len(df)} rows and columns: {list(df.columns)}")
                
                # Clean up column names and data
                df.columns = df.columns.str.strip()
                
                # Expected columns: "2025-01-27 00:00:00","2t","0 day fcst","surface","lon","lat","val"
                # Or similar format
                if 'lon' in df.columns and 'lat' in df.columns:
                    print(f"Lat range: {df['lat'].min():.3f} to {df['lat'].max():.3f}")
                    print(f"Lon range: {df['lon'].min():.3f} to {df['lon'].max():.3f}")
                    
                    if 'val' in df.columns:
                        print(f"Data range: {df['val'].min():.3f} to {df['val'].max():.3f}")
                        valid_data = df.dropna(subset=['lat', 'lon', 'val'])
                        print(f"Valid data points: {len(valid_data)}")
                        return valid_data
                
                return df
            else:
                print("CSV extraction failed or produced empty file")
                return None
                
        except subprocess.CalledProcessError as e:
            print(f"wgrib2 command failed: {e}")
            print(f"Error output: {e.stderr}")
            return None
        except Exception as e:
            print(f"wgrib2 extraction error: {e}")
            return None
    
    def extract_with_wgrib2_text(self, grib_file):
        """
        Extract using wgrib2 text output format
        """
        print("\n=== Method 2: wgrib2 text format ===")
        
        try:
            # Use -text option to get all grid point data
            text_file = os.path.join(self.temp_dir, 'arctic_data.txt')
            print(f"Extracting to text format: {text_file}")
            
            cmd_text = ['wgrib2', grib_file, '-text', text_file]
            result = subprocess.run(cmd_text, capture_output=True, text=True, check=True)
            
            if os.path.exists(text_file):
                # Read text data
                with open(text_file, 'r') as f:
                    lines = f.readlines()
                
                print(f"Read {len(lines)} lines from text file")
                
                # Parse the text data (format varies, typically one value per line)
                data_values = []
                for line in lines:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        try:
                            value = float(line)
                            data_values.append(value)
                        except ValueError:
                            continue
                
                print(f"Parsed {len(data_values)} data values")
                
                # Now get lat/lon coordinates
                return self._get_coordinates_for_data(grib_file, data_values)
            
        except Exception as e:
            print(f"Text extraction error: {e}")
            return None
    
    def _get_coordinates_for_data(self, grib_file, data_values):
        """
        Get lat/lon coordinates to match with data values
        """
        try:
            # Use wgrib2 to get lat/lon for each grid point
            latlon_file = os.path.join(self.temp_dir, 'latlon.txt')
            
            # Extract lat/lon using -lola option (custom grid)
            # First get grid dimensions
            cmd_nxny = ['wgrib2', grib_file, '-nxny']
            result = subprocess.run(cmd_nxny, capture_output=True, text=True, check=True)
            nxny_output = result.stdout.strip()
            print(f"Grid dimensions: {nxny_output}")
            
            # Parse nx, ny
            if '(' in nxny_output and ')' in nxny_output:
                coords_part = nxny_output.split('(')[1].split(')')[0]
                nx, ny = map(int, coords_part.split(' x '))
                print(f"Grid size: {nx} x {ny} = {nx*ny} points")
                
                # Create a lat/lon grid file using -lola option
                # This requires specifying output grid, but we want the native grid
                # Use -latlon option instead
                latlon_file = os.path.join(self.temp_dir, 'coords.txt')
                cmd_latlon = ['wgrib2', grib_file, '-latlon', latlon_file, '1', '1', '0']
                result = subprocess.run(cmd_latlon, capture_output=True, text=True, check=True)
                
                if os.path.exists(latlon_file):
                    with open(latlon_file, 'r') as f:
                        coord_lines = f.readlines()
                    
                    coordinates = []
                    for line in coord_lines:
                        parts = line.strip().split()
                        if len(parts) >= 2:
                            try:
                                lon, lat = float(parts[0]), float(parts[1])
                                coordinates.append((lat, lon))
                            except ValueError:
                                continue
                    
                    print(f"Extracted {len(coordinates)} coordinate pairs")
                    
                    # Match coordinates with data
                    min_len = min(len(coordinates), len(data_values))
                    
                    df = pd.DataFrame({
                        'latitude': [coord[0] for coord in coordinates[:min_len]],
                        'longitude': [coord[1] for coord in coordinates[:min_len]],
                        'value': data_values[:min_len]
                    })
                    
                    return df
            
        except Exception as e:
            print(f"Coordinate extraction error: {e}")
            return None
    
    def extract_with_pygrib_robust(self, grib_file):
        """
        Extract using pygrib with robust error handling
        """
        print("\n=== Method 3: pygrib robust extraction ===")
        
        try:
            import pygrib
            
            grbs = pygrib.open(grib_file)
            print(f"Opened GRIB file with {grbs.messages} messages")
            
            all_data = []
            
            for i, grb in enumerate(grbs):
                print(f"\nProcessing message {i+1}: {grb.name}")
                print(f"Grid type: {grb.gridType}")
                
                try:
                    # Try to get lat/lon coordinates
                    lats, lons = grb.latlons()
                    data = grb.values
                    
                    print(f"Successfully extracted coordinates and data")
                    print(f"Shape: {lats.shape}, Data range: {data.min():.3f} to {data.max():.3f}")
                    
                    # Flatten and create DataFrame
                    lats_flat = lats.flatten()
                    lons_flat = lons.flatten()
                    data_flat = data.flatten()
                    
                    # Remove invalid points
                    valid_mask = ~np.isnan(data_flat) & ~np.isnan(lats_flat) & ~np.isnan(lons_flat)
                    
                    valid_data = pd.DataFrame({
                        'latitude': lats_flat[valid_mask],
                        'longitude': lons_flat[valid_mask],
                        'value': data_flat[valid_mask],
                        'parameter': grb.name,
                        'message': i+1
                    })
                    
                    all_data.append(valid_data)
                    print(f"Added {len(valid_data)} valid points for {grb.name}")
                    
                except Exception as e:
                    print(f"Error processing message {i+1}: {e}")
                    
                    # Try alternative coordinate extraction
                    try:
                        print("Attempting alternative coordinate extraction...")
                        
                        # Get grid parameters
                        grid_params = {}
                        for attr in ['Ni', 'Nj', 'latitudeOfFirstGridPointInDegrees', 
                                   'longitudeOfFirstGridPointInDegrees', 'DxInMetres', 'DyInMetres']:
                            if hasattr(grb, attr):
                                grid_params[attr] = getattr(grb, attr)
                        
                        print(f"Grid parameters: {grid_params}")
                        
                        # Manual coordinate calculation using pyproj if available
                        coords_df = self._manual_polar_coordinates(grb, grid_params)
                        if coords_df is not None:
                            all_data.append(coords_df)
                            
                    except Exception as e2:
                        print(f"Alternative extraction also failed: {e2}")
                        continue
            
            grbs.close()
            
            if all_data:
                combined_df = pd.concat(all_data, ignore_index=True)
                print(f"\nCombined data: {len(combined_df)} total points")
                return combined_df
            else:
                return None
                
        except ImportError:
            print("pygrib not available. Install with: pip install pygrib")
            return None
        except Exception as e:
            print(f"pygrib extraction error: {e}")
            return None
    
    def _manual_polar_coordinates(self, grb, grid_params):
        """
        Manually calculate polar stereographic coordinates
        """
        try:
            if 'Ni' not in grid_params or 'Nj' not in grid_params:
                return None
            
            ni, nj = int(grid_params['Ni']), int(grid_params['Nj'])
            
            # Try with pyproj if available
            try:
                from pyproj import Proj, transform
                
                # Get projection parameters
                lat_0 = grid_params.get('latitudeOfFirstGridPointInDegrees', 90.0)
                lon_0 = grid_params.get('longitudeOfFirstGridPointInDegrees', 0.0)
                dx = grid_params.get('DxInMetres', 25000)
                dy = grid_params.get('DyInMetres', 25000)
                
                # Create coordinate grids
                x = np.arange(ni) * dx
                y = np.arange(nj) * dy
                X, Y = np.meshgrid(x, y)
                
                # Define projections
                proj_polar = Proj(proj='stere', lat_0=90, lon_0=lon_0, lat_ts=lat_0, ellps='sphere')
                proj_latlon = Proj(proj='latlong', ellps='sphere')
                
                # Transform to lat/lon
                lons, lats = transform(proj_polar, proj_latlon, X.flatten(), Y.flatten())
                
                # Get data values
                data = grb.values.flatten()
                
                # Create DataFrame
                df = pd.DataFrame({
                    'latitude': lats,
                    'longitude': lons,
                    'value': data,
                    'parameter': grb.name,
                    'method': 'manual_pyproj'
                })
                
                # Remove invalid points
                valid_mask = ~np.isnan(df['value']) & (df['latitude'] != 0) & (df['longitude'] != 0)
                df = df[valid_mask]
                
                print(f"Manual coordinate calculation: {len(df)} points")
                return df
                
            except ImportError:
                print("pyproj not available for coordinate transformation")
                return None
                
        except Exception as e:
            print(f"Manual coordinate calculation error: {e}")
            return None
    
    def extract_all_methods(self, grib_file):
        """
        Try all extraction methods and return the best result
        """
        print(f"\n🌊 Starting comprehensive Arctic GRIB extraction: {grib_file}")
        print("=" * 60)
        
        # Check tool availability
        tools = self.check_tools()
        
        results = []
        
        # Method 1: wgrib2 comprehensive (most reliable)
        if tools['wgrib2']:
            result1 = self.extract_with_wgrib2_comprehensive(grib_file)
            if result1 is not None and len(result1) > 0:
                results.append(('wgrib2_csv', result1))
                print(f"✓ wgrib2 CSV method: {len(result1)} points")
            else:
                # Try text method
                result1b = self.extract_with_wgrib2_text(grib_file)
                if result1b is not None and len(result1b) > 0:
                    results.append(('wgrib2_text', result1b))
                    print(f"✓ wgrib2 text method: {len(result1b)} points")
        
        # Method 2: pygrib
        if tools['pygrib']:
            result2 = self.extract_with_pygrib_robust(grib_file)
            if result2 is not None and len(result2) > 0:
                results.append(('pygrib', result2))
                print(f"✓ pygrib method: {len(result2)} points")
        
        # Return the best result (most data points)
        if results:
            best_method, best_data = max(results, key=lambda x: len(x[1]))
            print(f"\n🎯 Best result: {best_method} with {len(best_data)} data points")
            
            # Add some summary statistics
            if 'latitude' in best_data.columns:
                print(f"Latitude range: {best_data['latitude'].min():.3f}° to {best_data['latitude'].max():.3f}°")
            if 'longitude' in best_data.columns:
                print(f"Longitude range: {best_data['longitude'].min():.3f}° to {best_data['longitude'].max():.3f}°")
            if 'value' in best_data.columns or 'val' in best_data.columns:
                val_col = 'value' if 'value' in best_data.columns else 'val'
                print(f"Data range: {best_data[val_col].min():.3f} to {best_data[val_col].max():.3f}")
            
            return best_data, best_method
        else:
            print("❌ All extraction methods failed")
            return None, None
    
    def save_results(self, data, output_file, method_used):
        """
        Save extracted data to various formats
        """
        if data is None:
            print("No data to save")
            return
        
        base_name = os.path.splitext(output_file)[0]
        
        # Save as CSV
        csv_file = f"{base_name}.csv"
        data.to_csv(csv_file, index=False)
        print(f"Saved to CSV: {csv_file}")
        
        # Save as NetCDF if xarray available
        try:
            import xarray as xr
            
            # Convert to xarray dataset
            if 'latitude' in data.columns and 'longitude' in data.columns:
                ds = xr.Dataset({
                    'data': (['point'], data['value'] if 'value' in data.columns else data['val']),
                    'latitude': (['point'], data['latitude']),
                    'longitude': (['point'], data['longitude'])
                })
                
                ds.attrs['extraction_method'] = method_used
                ds.attrs['source'] = 'Arctic GRIB extraction'
                
                nc_file = f"{base_name}.nc"
                ds.to_netcdf(nc_file)
                print(f"Saved to NetCDF: {nc_file}")
                
        except ImportError:
            print("xarray not available, skipping NetCDF output")
        
        # Save metadata
        metadata = {
            'extraction_method': method_used,
            'total_points': len(data),
            'columns': list(data.columns),
            'latitude_range': [float(data['latitude'].min()), float(data['latitude'].max())] if 'latitude' in data.columns else None,
            'longitude_range': [float(data['longitude'].min()), float(data['longitude'].max())] if 'longitude' in data.columns else None,
        }
        
        metadata_file = f"{base_name}_metadata.json"
        with open(metadata_file, 'w') as f:
            json.dump(metadata, f, indent=2)
        print(f"Saved metadata: {metadata_file}")


def main():
    """
    Main function to extract Arctic GRIB data
    """
    # Example usage - replace with your actual file path
    grib_file = "/tmp/tmp0cvj_act.grib2"  # Your Arctic GRIB file
    
    if not os.path.exists(grib_file):
        print(f"GRIB file not found: {grib_file}")
        print("Please update the grib_file path in the script")
        return
    
    # Create extractor
    extractor = ArcticGRIBExtractor()
    
    # Extract all data
    data, method = extractor.extract_all_methods(grib_file)
    
    if data is not None:
        print(f"\n🎉 SUCCESS! Extracted {len(data)} data points using {method}")
        print("\nFirst 5 rows:")
        print(data.head())
        
        # Save results
        output_file = "arctic_extracted_data"
        extractor.save_results(data, output_file, method)
        
        print(f"\n📁 Results saved to:")
        print(f"  - {output_file}.csv")
        print(f"  - {output_file}_metadata.json")
        
    else:
        print("\n❌ Failed to extract data from Arctic GRIB file")
        print("Possible solutions:")
        print("1. Install wgrib2: brew install wgrib2 (macOS) or apt-get install wgrib2 (Linux)")
        print("2. Install pygrib: pip install pygrib")
        print("3. Install pyproj: pip install pyproj")


if __name__ == "__main__":
    main()