import pandas as pd
import geopandas as gpd
import json
import logging
from shapely.geometry import Point
from multiprocessing import Pool, cpu_count
from tqdm import tqdm
import numpy as np
from datetime import datetime
import requests
import time
import folium
from folium.plugins import HeatMap

class SolarAnalyzer:
    def __init__(self, api_key):
        # Set up logging
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)
        # Create console handler with a higher log level
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)
        # Create formatter and add it to the handlers
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        ch.setFormatter(formatter)
        # Add the handlers to the logger
        if not self.logger.handlers:
            self.logger.addHandler(ch)

        self.logger.info("Initializing SolarAnalyzer...")

        # Load the Gujarat GeoJSON data
        geojson_path = 'gujarat.geojson'  # Update the path as needed
        try:
            with open(geojson_path, 'r') as file:
                gujarat_geojson = json.load(file)
            self.gujarat_gdf = gpd.GeoDataFrame.from_features(gujarat_geojson['features'])
            self.gujarat_gdf.crs = "EPSG:4326"
            self.logger.info(f"Loaded Gujarat GeoJSON with {len(self.gujarat_gdf)} features.")
        except Exception as e:
            self.logger.error(f"Error loading GeoJSON file: {e}")
            raise e

        self.grid_resolution = 0.2  # Degrees
        self.grid_points = self._create_grid_points()
        self.land_usage_factor = 0.1  # Placeholder for usable land
        self.api_key = api_key  # NASA POWER API key

    def _create_grid_points(self) -> list:
        """Generate grid points within the Gujarat boundary."""
        self.logger.info("Generating grid points within Gujarat...")
        points = []

        minx, miny, maxx, maxy = self.gujarat_gdf.total_bounds
        latitudes = np.arange(miny, maxy + self.grid_resolution, self.grid_resolution)
        longitudes = np.arange(minx, maxx + self.grid_resolution, self.grid_resolution)

        # Precompute geometry for faster 'contains' checks
        gujarat_union = self.gujarat_gdf.unary_union

        for lat in tqdm(latitudes, desc="Generating Latitude Points"):
            for lon in longitudes:
                point = Point(lon, lat)
                if gujarat_union.contains(point):
                    points.append((lat, lon))

        self.logger.info(f"Generated {len(points)} grid points within Gujarat.")
        return points

    def fetch_solar_data(self, params):
        """
        Fetch solar data from NASA POWER API for given latitude and longitude.
        params: tuple containing (lat, lon)
        Returns climatology data along with lat and lon.
        """
        lat, lon = params
        lat = round(lat, 5)
        lon = round(lon, 5)
        self.logger.debug(f"Fetching data for Lat: {lat}, Lon: {lon}")

        # NASA POWER API endpoint
        url = "https://power.larc.nasa.gov/api/temporal/daily/point"

        # Define the parameters for the API request
        payload = {
            'parameters': ['ALLSKY_SFC_SW_DWN', 'SWDIR_SFC', 'CLRSKY_SFC_SW_DWN', 'T2M', 'WS10M'],
            'community': 'SB',      # Solar and Buildings community
            'longitude': lon,
            'latitude': lat,
            'start': '20200101',    # Start date in YYYYMMDD
            'end': '20201231',      # End date in YYYYMMDD
            'format': 'JSON',
            'user': self.api_key
        }

        try:
            response = requests.get(url, params=payload, timeout=30)
            response.raise_for_status()
            data = response.json()

            # Extract climatology data
            climatology = data.get('properties', {}).get('parameter', {})
            if not climatology:
                self.logger.warning(f"No climatology data for Lat: {lat}, Lon: {lon}")
                return None

            return (climatology, lat, lon)

        except requests.exceptions.RequestException as e:
            self.logger.error(f"Request error for Lat: {lat}, Lon: {lon} -> {e}")
            return None
        except json.JSONDecodeError:
            self.logger.error(f"JSON decode error for Lat: {lat}, Lon: {lon}")
            return None
        except Exception as e:
            self.logger.error(f"Unexpected error for Lat: {lat}, Lon: {lon} -> {e}")
            return None

    def collect_data(self, start_year: int, end_year: int) -> pd.DataFrame:
        """Fetch solar and meteorological data for all grid points."""
        self.logger.info("Starting data collection from NASA POWER API...")
        pool = Pool(processes=cpu_count())
        params = [(lat, lon) for lat, lon in self.grid_points]

        results = []
        try:
            for result in tqdm(pool.imap_unordered(self.fetch_solar_data, params), total=len(params), desc="Fetching Data"):
                results.append(result)
        except KeyboardInterrupt:
            self.logger.warning("Process interrupted by user.")
            pool.terminate()
            pool.join()
            raise
        except Exception as e:
            self.logger.error(f"Error during data collection: {e}")
            pool.terminate()
            pool.join()
            raise

        pool.close()
        pool.join()

        all_records = []
        for result in results:
            if result:
                climatology, lat, lon = result
                # Aggregate daily data into monthly averages
                monthly_data = {}
                for date_str, values in climatology.items():
                    if len(date_str) != 8:
                        continue  # Skip unexpected keys
                    month = int(date_str[4:6])
                    if month not in monthly_data:
                        monthly_data[month] = {
                            'ALLSKY_SFC_SW_DWN': [],
                            'SWDIR_SFC': [],
                            'CLRSKY_SFC_SW_DWN': [],
                            'T2M': [],
                            'WS10M': []
                        }
                    monthly_data[month]['ALLSKY_SFC_SW_DWN'].append(values.get('ALLSKY_SFC_SW_DWN'))
                    monthly_data[month]['SWDIR_SFC'].append(values.get('SWDIR_SFC'))
                    monthly_data[month]['CLRSKY_SFC_SW_DWN'].append(values.get('CLRSKY_SFC_SW_DWN'))
                    monthly_data[month]['T2M'].append(values.get('T2M'))
                    monthly_data[month]['WS10M'].append(values.get('WS10M'))

                # Compute monthly averages
                for month, data in monthly_data.items():
                    record = {
                        'month': month,
                        'latitude': lat,
                        'longitude': lon,
                        'ALLSKY_SFC_SW_DWN_avg': np.nanmean(data['ALLSKY_SFC_SW_DWN']),
                        'SWDIR_SFC_avg': np.nanmean(data['SWDIR_SFC']),
                        'CLRSKY_SFC_SW_DWN_avg': np.nanmean(data['CLRSKY_SFC_SW_DWN']),
                        'T2M_avg': np.nanmean(data['T2M']),
                        'WS10M_avg': np.nanmean(data['WS10M']),
                    }
                    all_records.append(record)

        self.logger.info(f"Collected data for {len(all_records)} records.")
        return pd.DataFrame(all_records)

    def calculate_solar_power(self, df: pd.DataFrame) -> pd.Series:
        """Compute the solar power potential based on meteorological data."""
        panel_efficiency = 0.20
        system_losses = 0.14
        temperature_coefficient = -0.0035  # Not used in current calculation

        # Convert ALLSKY_SFC_SW_DWN from W/m² to kWh/m²/day
        # Assuming data is daily total in Wh/m²/day (verify with documentation)
        # Here, sum of daily total might be needed

        # For illustration, let's assume ALLSKY_SFC_SW_DWN_avg is kWh/m²/month
        # Convert to kWh/m²/day by dividing by average days per month (~30)
        solar_radiation_kwh_m2_day = df['ALLSKY_SFC_SW_DWN_avg'] / 30  # Simple approximation

        # Calculate power output in kWh/m²/day
        power_output_kwh_m2_day = (
            solar_radiation_kwh_m2_day *
            panel_efficiency *
            (1 - system_losses)
        )

        # Convert kWh/m²/day to MW/km²/year
        # 1 kWh/m²/day * 365 = 365 kWh/m²/year
        # 1 kWh/m²/year = 0.001 MWh/m²/year
        # 1 km² = 1,000,000 m²
        power_MW_km2_year = power_output_kwh_m2_day * 365 * 0.001 * self.land_usage_factor * 1e6 / 1e6  # Simplifies to kWh/m²/year * 0.001 * land_usage_factor

        return power_MW_km2_year

    def create_heatmap(self, df: pd.DataFrame, output_path: str = "solar_heatmap.html"):
        """Generate an interactive heatmap of solar power potential over Gujarat."""
        self.logger.info("Creating heatmap...")

        # Aggregate data by location (average over months)
        grid_summary = df.groupby(['latitude', 'longitude'])['total_potential'].mean().reset_index()

        # Prepare data for heatmap
        heat_data = [
            [row['latitude'], row['longitude'], row['total_potential']]
            for index, row in grid_summary.iterrows()
        ]

        # Initialize Folium map centered on Gujarat
        if grid_summary.empty:
            self.logger.error("No data available to create heatmap.")
            return

        avg_lat = grid_summary['latitude'].mean()
        avg_lon = grid_summary['longitude'].mean()
        m = folium.Map(location=[avg_lat, avg_lon], zoom_start=7, tiles='CartoDB positron')

        # Add HeatMap
        HeatMap(
            heat_data,
            min_opacity=0.5,
            max_val=grid_summary['total_potential'].max(),
            radius=15,
            blur=25,
            max_zoom=1,
            gradient={0.2: 'blue', 0.4: 'lime', 0.6: 'yellow', 0.8: 'orange', 1.0: 'red'}
        ).add_to(m)

        # Save to HTML
        m.save(output_path)
        self.logger.info(f"Heatmap saved to {output_path}")

    def analyze(self, start_year: int, end_year: int):
        """Execute the entire analysis pipeline."""
        self.logger.info(f"Starting analysis for years {start_year}-{end_year}...")
        # Step 1: Collect data
        data_df = self.collect_data(start_year, end_year)

        if data_df.empty:
            self.logger.error("No data collected. Exiting analysis.")
            return

        # Step 2: Calculate solar potential
        self.logger.info("Calculating solar power potential...")
        data_df['total_potential'] = self.calculate_solar_power(data_df)

        # Remove rows with NaN potentials
        data_df = data_df.dropna(subset=['total_potential'])

        # Step 3: Visualize the results
        self.create_heatmap(data_df)

        # Step 4: Save results to a CSV (optional)
        output_csv = f"solar_power_potential_{start_year}_{end_year}.csv"
        data_df.to_csv(output_csv, index=False)
        self.logger.info(f"Results saved to {output_csv}")

def main():
    # Replace with your actual NASA POWER API key
    nasa_power_api_key = 'YOUR_NASA_POWER_API_KEY'

    # Initialize the solar analyzer
    analyzer = SolarAnalyzer(api_key=nasa_power_api_key)

    # Define the date range for data collection (e.g., year 2020)
    start_year = 2020
    end_year = 2020  # Single year; modify as needed for multiple years

    # Execute the analysis
    analyzer.analyze(start_year, end_year)

if __name__ == "__main__":
        main()

This code is currently fetching exactly 0 data from NASA API, indicating that they don't have data for such precise coordinates. Can you fix this?