import numpy as np
import pandas as pd
import geopandas as gpd
import folium
import requests
from shapely.geometry import Point
from scipy.interpolate import griddata
import time
import branca.colormap as cm
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
import json
import shelve
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import logging
import threading

# --- 1. Set Up Logging ---
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.FileHandler("solar_map.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# --- 2. Rate Limiter Implementation ---
class RateLimiter:
    """
    Simple thread-safe rate limiter.
    Allows up to `max_calls` within `period` seconds.
    """
    def __init__(self, max_calls, period):
        self.max_calls = max_calls
        self.period = period
        self.lock = threading.Lock()
        self.calls = []
    
    def acquire(self):
        with self.lock:
            current = time.time()
            # Remove calls that are outside the current period
            while self.calls and self.calls[0] <= current - self.period:
                self.calls.pop(0)
            if len(self.calls) >= self.max_calls:
                # Sleep until the earliest call is outside the period
                sleep_time = self.period - (current - self.calls[0]) + 0.1  # Add a small buffer
                logger.info(f"Rate limit reached. Sleeping for {sleep_time:.2f} seconds.")
                time.sleep(sleep_time)
            self.calls.append(time.time())

# Initialize RateLimiter (e.g., 1 call per 2 seconds)
rate_limiter = RateLimiter(max_calls=1, period=2)

# --- 3. Session with Retry Strategy ---
def create_session_with_retry():
    """
    Create a requests session with retry strategy
    """
    session = requests.Session()
    # Configure retry strategy
    retries = Retry(
        total=10,  # Increased from 5 to 10 retries
        backoff_factor=2,  # Increased backoff_factor for exponential backoff
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["HEAD", "GET", "OPTIONS"],  # Updated parameter name
        respect_retry_after_header=True
    )
    adapter = HTTPAdapter(max_retries=retries)
    session.mount('https://', adapter)
    session.mount('http://', adapter)
    return session

# --- 4. Fetch Data with Adaptive Rate Limiting ---
def fetch_nasa_power_data(lat, lon, cache, session):
    """
    Fetch solar radiation data from NASA POWER API with caching and adaptive rate limiting.
    """
    key = f"{lat}_{lon}"
    if key in cache:
        logger.debug(f"Cache hit for lat={lat}, lon={lon}")
        return cache[key]
    
    base_url = "https://power.larc.nasa.gov/api/temporal/daily/point"
    params = {
        'parameters': 'ALLSKY_SFC_SW_DWN',
        'community': 'RE',
        'longitude': str(lon),
        'latitude': str(lat),
        'start': '20230101',
        'end': '20231231',
        'format': 'JSON'
    }
    
    while True:
        try:
            rate_limiter.acquire()  # Ensure rate limiting
            logger.debug(f"Sending request for lat={lat}, lon={lon}")
            response = session.get(base_url, params=params)
            if response.status_code == 429:
                # Handle rate limiting
                retry_after = response.headers.get('Retry-After')
                if retry_after:
                    sleep_time = int(retry_after) + 1
                    logger.warning(f"Received 429 for lat={lat}, lon={lon}. Retrying after {sleep_time} seconds.")
                    time.sleep(sleep_time)
                else:
                    # Default wait time if Retry-After not provided
                    sleep_time = 60
                    logger.warning(f"Received 429 for lat={lat}, lon={lon}. Retrying after {sleep_time} seconds.")
                    time.sleep(sleep_time)
                continue  # Retry the request after sleeping
            response.raise_for_status()
            data = response.json()
            cache[key] = data
            logger.debug(f"Successfully fetched data for lat={lat}, lon={lon}")
            return data
        except requests.exceptions.RequestException as e:
            logger.error(f"RequestException for lat={lat}, lon={lon}: {e}. Retrying in 10 seconds.")
            time.sleep(10)
        except json.JSONDecodeError as e:
            logger.error(f"JSONDecodeError for lat={lat}, lon={lon}: {e}. Skipping this point.")
            cache[key] = None
            return None

# --- 5. Calculate Solar Potential ---
def calculate_solar_potential(power_data, panel_efficiency=0.2):
    """
    Calculate solar potential from radiation data.
    """
    if not power_data:
        return None
    try:
        daily_data = power_data['properties']['parameter']['ALLSKY_SFC_SW_DWN']
        radiation_values = [v * 0.277778 for v in daily_data.values() 
                            if isinstance(v, (int, float)) and v != -999]
        if not radiation_values:
            return None
        return np.mean(radiation_values) * 365 * panel_efficiency
    except KeyError as e:
        logger.error(f"KeyError while calculating solar potential: {e}")
        return None
    except Exception as e:
        logger.error(f"Unexpected error while calculating solar potential: {e}")
        return None

# --- 6. Helper Function for Fetching and Calculating ---
def fetch_and_calculate(lat, lon, cache, session):
    """
    Helper function to fetch power data and calculate solar potential.
    """
    power_data = fetch_nasa_power_data(lat, lon, cache, session)
    potential = calculate_solar_potential(power_data)
    if potential is not None:
        return {'latitude': lat, 'longitude': lon, 'potential': potential}
    return None

# --- 7. Create Solar Map Function ---
def create_india_solar_map(geojson_path='india-soi.geojson'):
    # Read India GeoJSON
    logger.info(f"Reading GeoJSON file from: {geojson_path}")
    india = gpd.read_file(geojson_path)
    
    # Simplify geometry for performance
    india_simplified = india.geometry.simplify(0.1)
    
    # Combine all geometries into a single union using union_all()
    india_union = india_simplified.union_all()
    
    # Get India's bounds (minx, miny, maxx, maxy)
    bounds = india_union.bounds
    logger.info(f"Map bounds: {bounds}")
    
    # Create base Folium map
    m = folium.Map(location=[20.5937, 78.9629], zoom_start=5, tiles='CartoDB positron')
    
    # Define grid parameters (finer grid)
    lat_step = 0.5  # Decreased step size for higher precision
    lon_step = 0.5
    lat_range = np.arange(bounds[1], bounds[3] + lat_step, lat_step)
    lon_range = np.arange(bounds[0], bounds[2] + lon_step, lon_step)
    
    # Initialize cache
    cache_file = 'nasa_power_cache.db'
    cache = shelve.open(cache_file)
    
    # Initialize session with retry strategy
    session = create_session_with_retry()
    
    # Define number of workers (adaptive based on rate limits)
    max_workers = 3  # Reduced from 5 to 3 to minimize concurrent requests
    
    logger.info(f"Starting data collection with {max_workers} workers...")
    
    solar_data = []
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit tasks
        future_to_point = {
            executor.submit(fetch_and_calculate, lat, lon, cache, session): (lat, lon)
            for lat in lat_range
            for lon in lon_range
            if india_union.contains(Point(lon, lat))
        }
        
        total_grid_points = len(future_to_point)
        completed = 0
        for future in as_completed(future_to_point):
            lat, lon = future_to_point[future]
            try:
                result = future.result()
                if result:
                    solar_data.append(result)
            except Exception as e:
                logger.error(f"Error processing point at lat={lat}, lon={lon}: {e}")
            finally:
                completed += 1
                if completed % 50 == 0 or completed == total_grid_points:
                    logger.info(f"Processed {completed}/{total_grid_points} points")
    
    cache.close()
    
    solar_df = pd.DataFrame(solar_data)
    
    if solar_df.empty:
        raise ValueError("No solar data collected")
    
    logger.info("Creating interpolation grid...")
    
    # Create interpolation grid (finer resolution)
    grid_size = 400  # Increased grid size for higher resolution
    grid_lat = np.linspace(bounds[1], bounds[3], grid_size)
    grid_lon = np.linspace(bounds[0], bounds[2], grid_size)
    grid_lon, grid_lat = np.meshgrid(grid_lon, grid_lat)
    
    # Interpolate data
    points = solar_df[['longitude', 'latitude']].values
    values = solar_df['potential'].values
    grid_z = griddata(points, values, (grid_lon, grid_lat), method='linear')
    
    # Create mask for points within India's boundary
    logger.info("Creating mask for India's boundary...")
    grid_points = [Point(lon, lat) for lon, lat in zip(grid_lon.flatten(), grid_lat.flatten())]
    india_gdf = gpd.GeoSeries(grid_points)
    mask = india_gdf.within(india_union).values.reshape(grid_z.shape)
    
    # Apply mask
    grid_z = np.ma.masked_array(grid_z, ~mask)
    
    logger.info("Creating visualization...")
    
    # Define expanded color scheme using 'viridis' equivalents
    colormap = cm.LinearColormap(
        ['#440154', '#482878', '#3E4989', '#31688E', '#26828E',
         '#1F9E89', '#35B779', '#6DCD59', '#B4DE2C', '#FDE725'],
        vmin=np.nanmin(values),
        vmax=np.nanmax(values),
        caption='Solar Potential (kWh/m²/year)'
    )
    
    # Normalize and map data to colors
    norm = plt.Normalize(vmin=colormap.vmin, vmax=colormap.vmax)
    cmap = ListedColormap(colormap.colors)
    
    # Generate high-resolution image
    fig, ax = plt.subplots(figsize=(16, 12))  # Increased figure size
    ax.set_axis_off()
    ax.imshow(
        grid_z, 
        cmap=cmap, 
        extent=(bounds[0], bounds[2], bounds[1], bounds[3]),
        origin='lower', 
        aspect='auto'
    )
    
    # Save the image with high DPI
    image_path = 'solar_potential_high_res.png'
    plt.savefig(image_path, bbox_inches='tight', pad_inches=0, dpi=600)
    plt.close(fig)
    
    # Verify image creation
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"Failed to create image at {image_path}")
    
    # Add ImageOverlay to Folium map
    folium.raster_layers.ImageOverlay(
        name='Solar Potential',
        image=image_path,
        bounds=[[bounds[1], bounds[0]], [bounds[3], bounds[2]]],
        opacity=0.6,
        interactive=True,
        cross_origin=False,
        zindex=1,
    ).add_to(m)
    
    # Add India boundary
    folium.GeoJson(
        india,
        style_function=lambda x: {
            'fillColor': 'none',
            'color': 'black',
            'weight': 1.5,
            'fillOpacity': 0
        },
        name='India Boundary'
    ).add_to(m)
    
    # Add colormap legend
    colormap.add_to(m)
    
    # Add layer control
    folium.LayerControl().add_to(m)
    
    # Save data
    output_csv = 'india_solar_data.csv'
    output_html = 'india_solar_potential.html'
    
    logger.info(f"Saving data to {output_csv}")
    solar_df.to_csv(output_csv, index=False)
    
    logger.info(f"Saving map to {output_html}")
    m.save(output_html)
    
    # Optional: Remove the high-resolution image to save space
    # os.remove(image_path)
    
    logger.info("Solar potential map generation completed successfully.")
    
    return m, solar_df

# --- 8. Main Function ---
def main():
    """
    Main function to run the application.
    """
    try:
        logger.info("Starting solar potential map generation with adaptive rate limiting...")
        solar_map, solar_data = create_india_solar_map()
        logger.info("Process completed successfully!")
        logger.info("Map has been saved as 'india_solar_potential.html'")
        logger.info("Raw data has been saved as 'india_solar_data.csv'")
    except Exception as e:
        logger.error(f"An error occurred: {e}")

# --- 9. Entry Point ---
if __name__ == "__main__":
    main()

The above code needs following improvements
1. The code fails to take in data for the entire north India and only shows output on South India
2. The overlay generated is not 1:1 fit for the geojson
3. There's no progress indicator
4. Comments must be removed
5. The following message is the only output when running on Colab:
/usr/local/lib/python3.10/dist-packages/shapely/decorators.py:79: RuntimeWarning: coroutine 'main' was never awaited
  for arr, old_flag in zip(array_args, old_flags):
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
/usr/local/lib/python3.10/dist-packages/shapely/decorators.py:79: RuntimeWarning: coroutine 'create_india_solar_map' was never awaited
  for arr, old_flag in zip(array_args, old_flags):
RuntimeWarning: Enable tracemalloc to get the object allocation traceback