import os
import json
import imdlib as imd
import numpy as np
import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

CITIES = {
    "Mumbai": {"lat": 18.9878, "lon": 72.8364},
    "Pune": {"lat": 18.5204, "lon": 73.8567},
    "Nagpur": {"lat": 21.1458, "lon": 79.0882},
    "Nashik": {"lat": 19.9975, "lon": 73.7898},
    "Chhatrapati Sambhajinagar": {"lat": 19.8762, "lon": 75.3433},
    "Panaji": {"lat": 15.4909, "lon": 73.8278},
    "Margao": {"lat": 15.2736, "lon": 73.9580},
    "Vasco da Gama": {"lat": 15.3973, "lon": 73.8122},
    "Mapusa": {"lat": 15.5937, "lon": 73.8105},
    "Ponda": {"lat": 15.4026, "lon": 74.0156},
    "Bicholim": {"lat": 15.5908, "lon": 73.9525},
    "Curchorem": {"lat": 15.2530, "lon": 74.1165},
    "Sanguem": {"lat": 15.2289, "lon": 74.1539},
    "Canacona": {"lat": 15.0118, "lon": 74.0223},
    "Pernem": {"lat": 15.7187, "lon": 73.7997},
    "Bengaluru": {"lat": 12.9716, "lon": 77.5946},
    "Mysuru": {"lat": 12.2958, "lon": 76.6394},
    "Hubballi": {"lat": 15.3647, "lon": 75.1240},
    "Mangaluru": {"lat": 12.9141, "lon": 74.8560},
    "Belagavi": {"lat": 15.8497, "lon": 74.4977},
    "Kalaburagi": {"lat": 17.3297, "lon": 76.8343},
    "Davanagere": {"lat": 14.4644, "lon": 75.9218},
    "Ballari": {"lat": 15.1394, "lon": 76.9214},
    "Vijayapura": {"lat": 16.8302, "lon": 75.7100},
    "Shivamogga": {"lat": 13.9299, "lon": 75.5681},
    "Tumakuru": {"lat": 13.3392, "lon": 77.1010},
    "Raichur": {"lat": 16.2076, "lon": 77.3463},
    "Bidar": {"lat": 17.9104, "lon": 77.5199},
    "Hosapete": {"lat": 15.2711, "lon": 76.3908},
    "Gadag": {"lat": 15.4297, "lon": 75.6322},
    "Hassan": {"lat": 13.0033, "lon": 76.1004},
    "Udupi": {"lat": 13.3409, "lon": 74.7421},
    "Chitradurga": {"lat": 14.2251, "lon": 76.3980},
    "Kolar": {"lat": 13.1367, "lon": 78.1291},
    "Mandya": {"lat": 12.5218, "lon": 76.8951}
}

def fetch_and_process():
    # 2014 to 2023 (10 complete years)
    start_yr = 2014
    end_yr = 2023
    file_dir = './imd_raw_data'
    os.makedirs(file_dir, exist_ok=True)
    
    variables = ['tmax', 'tmin', 'rain']
    datasets = {}
    
    for var in variables:
        print(f"Downloading {var} from {start_yr} to {end_yr}...")
        try:
            imd.get_data(var, start_yr, end_yr, fn_format='yearwise', file_dir=file_dir)
            data = imd.open_data(var, start_yr, end_yr, 'yearwise', file_dir)
            datasets[var] = data.get_xarray()
            print(f"Successfully loaded {var} dataset.")
        except Exception as e:
            print(f"Error processing {var}: {e}")
            return

    # Extract historical extremes for each city
    results = {}
    
    print("Processing historical data for cities...")
    for city_name, coords in CITIES.items():
        lat = coords['lat']
        lon = coords['lon']
        
        city_data = {}
        try:
            for var in variables:
                ds = datasets[var]
                # Nearest grid point
                ts = ds[var].sel(lat=lat, lon=lon, method='nearest')
                
                # We replace 99.9 (IMD nodata) with NaN
                ts = ts.where(ts < 99.9)
                
                if var == 'tmax':
                    city_data['all_time_max'] = float(ts.max().values)
                    city_data['avg_max'] = float(ts.mean().values)
                elif var == 'tmin':
                    city_data['all_time_min'] = float(ts.min().values)
                elif var == 'rain':
                    city_data['max_daily_rain'] = float(ts.max().values)
                    
        except Exception as e:
            print(f"Error calculating stats for {city_name}: {e}")
            
        results[city_name] = city_data

    # Save to frontend data directory
    out_dir = os.path.join(os.path.dirname(__file__), '..', 'frontend', 'src', 'data')
    os.makedirs(out_dir, exist_ok=True)
    out_file = os.path.join(out_dir, 'imd_historical_baseline.json')
    
    with open(out_file, 'w') as f:
        json.dump(results, f, indent=2)
        
    print(f"Successfully saved IMD Historical Baseline to {out_file}")

if __name__ == "__main__":
    fetch_and_process()