import boto3
import pyart
import numpy as np
import cmweather
import matplotlib.pyplot as plt
import torch
import os
import cartopy.crs as ccrs
import dask.bag as db
import tempfile

from glob import glob
from datetime import datetime, timedelta 
from botocore import UNSIGNED
from torchvision.io import decode_image
from torchvision import transforms
from botocore.config import Config

class RadarImage(object):
    """
    This class contains the basic data for predicting the lake-breeze front location
    from radar. It includes parameters for storing the radar data for later analysis
    in Py-ART and for inference in the lake-breeze prediction model.

    Parameters
    ----------
    pyart_object: :py:meth:`pyart.core.Radar` or str
        The PyART radar object that stores the radar data. This could also be a link
        to the radar scan file (useful for batch processing to preserve memory).
    lat_range: 2-tuple
        The minimum and maximum latitude of the inference domain.
    lon_range: 2-tuple
        The minimum and maximum longitude of the inference domain.
    grid_lat: ndarray
        The latitude of each point in the inference domain.
    grid_lon: ndarray
        The longitude of each point in the inference domain.
    pytorch_image: :py:meth:`torch.Tensor`
        The tensor containing the preprocessed radar scan for inference.
    lakebreeze_mask: 256 x 256 ndarray
        The inferred lake breeze mask, where 1 = lakebreeze and 0 = not a lake breeze. 
    times: list of np.datetime64('s')
        The epoch time of the radar scans.
    """
    pyart_object = None
    lat_range = None
    lon_range = None
    grid_lat = None
    grid_lon = None
    pytorch_image = None
    lakebreeze_mask = None
    aggregated_mask = None
    times = None
    
    def __getitem__(self, key):
        """
        Allows for indexing into the RadarImage object to get the lake breeze mask
        for a specific time.
        
        Parameters
        ----------
        key: int or np.datetime64('s')
            The index or time to retrieve the lake breeze mask for.

        Returns
        -------
        mask: ndarray
            The lake breeze mask for the specified time.
        """
        if len(self.lakebreeze_mask.shape) == 2:
            return self.lakebreeze_mask
        if isinstance(key, int):
            return self.lakebreeze_mask[key]
        elif isinstance(key, np.datetime64):
            return self.aggregate(start_time=key, end_time=key)
        elif isinstance(key, str):
            key = np.datetime64(key)
            return self.aggregate(start_time=key, end_time=key)
        elif isinstance(key, slice):
            return self.lakebreeze_mask[key]
        elif isinstance(key, list) or isinstance(key, np.ndarray):
            if all(isinstance(x, int) for x in key):
                return [self.lakebreeze_mask[x] for x in key]
            elif all(isinstance(x, (str, np.datetime64)) for x in key):
                masks = []
                for time in key:
                    if isinstance(time, str):
                        time = np.datetime64(time)
                    masks.append(self.aggregate(start_time=time, end_time=time))
                return masks
            else:
                raise TypeError("All keys in the list must be of the same type: int, str, or np.datetime64('s').")
        else:
            raise TypeError("Key must be an int, np.datetime64('s'), str, or slice.")
        
    def aggregate(self, start_time=None, end_time=None):
        """
        This function aggregates the lake breeze mask over a specified time period.
        If no time period is specified, it returns the sum of all of the masks.

        Parameters
        ----------
        start_time: str or np.datetime64('s')
            The start time for aggregation.
        end_time: str np.datetime64('s')
            The end time for aggregation.

        Returns
        -------
        aggregated_mask: RadarImage
            The aggregated lake breeze mask.
        """
        if start_time is None and end_time is None:
            self.aggregated_mask = np.sum(self.lakebreeze_mask, axis=0)
            return self.aggregated_mask

        mask = self.lakebreeze_mask.copy()
        if isinstance(start_time, str):
            start_time = np.datetime64(start_time)
        if isinstance(end_time, str):
            end_time = np.datetime64(end_time)
        if (start_time is None) ^ (end_time is None):
            raise ValueError("Both start_time and end_time must be specified for aggregation.")
        indices = np.where((self.times >= start_time) & (self.times <= end_time))[0]
        if len(indices) == 0:
            raise ValueError("No data available for the specified time range.")
        mask = mask[indices]
        self.aggregated_mask = np.sum(mask, axis=0)
        return self.aggregated_mask


def preprocess_radar_image(radar, rad_time=None, lat_range=(41.1280, 42.5680),
                           lon_range=(-88.7176, -87.2873),
                           bucket_name='unidata-nexrad-level2'):
    """
    This module will preprocess the NEXRAD radar data for inference into the lake-breeze
    prediction model of ADAM.

    Parameters
    ----------
    radar: str or :py:meth:`pyart.core.radar` object
        The 4-letter code for the radar to obtain the scan from. For Chicago, use KLOT.
    rad_time: ISO-format datestring
        The date/time string in YYYY-MM-DDTHH:MM:SS format for the radar scan. If None, then ADAM will
        get the latest scan. This is not used if radar is a :py:meth:`pyart.core.radar` object or string.
    lat_range: 2-tuple of floats
        The minimum and maximum latitude of the domain in degrees. Default is a centered
        domain around the KLOT Chicago area radar.
    lon_range: 2-tuple of floats
        The minimum and maximum longitude of the domain in degrees. Default is a centered
        domain around the KLOT Chicago area radar.
    bucket_name: str
        The NEXRAD S3 bucket to use. Default is 'unidata-nexrad-level2'.

    Returns
    -------
    image: :py:meth:`adam.io.RadarImage`
        The :py:meth:`RadarImage` object containing the radar scan, pre-processed image,
        and grid.
    """
    if isinstance(radar, str):
        if rad_time is None:
            right_now = datetime.utcnow()
        else:
            right_now = datetime.strptime(rad_time, "%Y-%m-%dT%H:%M:%S")
        yesterday = right_now - timedelta(days=1)
        year = right_now.year
        month = right_now.month
        day = right_now.day

        s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
        bucket_name = 'unidata-nexrad-level2'
        radar = "KLOT"
        prefix = f'{year}/{month:02d}/{day:02d}/{radar}'
        response = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
        file_list = [x['Key'] for x in response['Contents']]
        
        # Find yesterday's scans
        prefix = f'{yesterday.year}/{yesterday.month:02d}/{yesterday.day:02d}/{radar}'
        response = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
        file_list = file_list + [x['Key'] for x in response['Contents']]
        time_list = []
        for filepath in file_list:
            name = filepath.split("/")[-1]
            if name[-3:] == "MDM":
                time_list.append(
                    datetime.strptime(name, f"{radar}%Y%m%d_%H%M%S_V06_MDM"))
            else:
                time_list.append(
                    datetime.strptime(name, f"{radar}%Y%m%d_%H%M%S_V06"))

        time_list = np.array(time_list)
        path = f"s3://{bucket_name}/" + file_list[np.argmin(np.abs(time_list - right_now))]
        cur_radar = pyart.io.read_nexrad_archive(path)
    elif isinstance(radar, pyart.core.Radar):
        cur_radar = radar
    elif isinstance(radar, str):
        cur_radar = pyart.io.read(str)
    else:
        raise ValueError("The radar input must be a string or a PyART radar object.")

    disp = pyart.graph.RadarMapDisplay(cur_radar)
    fig, ax = plt.subplots(1, 1, figsize=(2.56, 2.56),
            subplot_kw=dict(projection=ccrs.PlateCarree(), frameon=False))

    disp.plot_ppi_map('reflectivity', sweep=0, min_lon=lon_range[0],
            ax=ax, max_lon=lon_range[1], min_lat=lat_range[0], max_lat=lat_range[1],
            embellish=False, vmin=-20, vmax=60, cmap='HomeyerRainbow',
            add_grid_lines=False, colorbar_flag=False, title_flag=False)
    ax.set_axis_off()
    fig.tight_layout(pad=0, w_pad=0, h_pad=0)
    with tempfile.NamedTemporaryFile(mode='w+b') as temp_file:
        fig.savefig(temp_file, dpi=100)
        plt.close(fig)
        # Transform image
        image = decode_image(temp_file.name)
    image = image[:3, :, :].float()
    transform = transforms.Compose([
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    image = torch.stack([transform(image)])
    lats = np.linspace(lat_range[1], lat_range[0], image.shape[3])
    lons = np.linspace(lon_range[0], lon_range[1], image.shape[2])

    rad_image = RadarImage()
    rad_image.pytorch_image = image
    rad_image.lat_range = lat_range
    rad_image.lon_range = lon_range
    rad_image.pyart_object = cur_radar
    rad_image.grid_lat = lats
    rad_image.grid_lon = lons
    center_lat = (rad_image.lat_range[0] + rad_image.lat_range[1]) / 2.
    center_lon = (rad_image.lon_range[0] + rad_image.lon_range[1]) / 2.
    rad_image.grid_x, rad_image.grid_y = _latlon_to_xy(
        rad_image.grid_lat, rad_image.grid_lon, center_lat, center_lon)
    rad_image.times = [np.datetime64(cur_radar.time["units"].split()[2])]
    
    return rad_image

def preprocess_radar_image_batch(file, lat_range=(41.1280, 42.5680),
                           lon_range=(-88.7176, -87.2873), parallel=False):
    """
    This module will preprocess the NEXRAD radar data for inference into the lake-breeze
    prediction model of ADAM.

    Parameters
    ----------
    radar: str or list
        The 4-letter code for the radar to obtain the scan from. For Chicago, use KLOT.
    lat_range: 2-tuple of floats
        The minimum and maximum latitude of the domain in degrees. Default is a centered
        domain around the KLOT Chicago area radar.
    lon_range: 2-tuple of floats
        The minimum and maximum longitude of the domain in degrees. Default is a centered
        domain around the KLOT Chicago area radar.
    parallel: bool
        If true, enable parallel preprocessing for large radar datasets using Dask.

    Returns
    -------
    image: :py:meth:`adam.io.RadarImage`
        The :py:meth:`RadarImage` object containing the radar scan, pre-processed image,
        and grid.
    """
    
    if isinstance(file, str):
        files = sorted(glob(file))
    else:
        files = file
    _pprocess = lambda x: _preprocess(x, lat_range, lon_range)

    if parallel:
        arr = db.from_sequence(files).map(_pprocess).compute()
        images = [x[0] for x in arr]
        times = [x[1] for x in arr]
    else:
        arr = map(_pprocess, files)
        images = [x[0] for x in arr]
        times = [x[1] for x in arr]

    images = torch.concat(images, axis=0)
    lats = np.linspace(lat_range[1], lat_range[0], images.shape[3])
    lons = np.linspace(lon_range[0], lon_range[1], images.shape[2])
    rad_image = RadarImage()
    rad_image.pytorch_image = images
    rad_image.lat_range = lat_range
    rad_image.lon_range = lon_range
    rad_image.pyart_object = files
    rad_image.grid_lat = lats
    rad_image.grid_lon = lons
    center_lat = (rad_image.lat_range[0] + rad_image.lat_range[1]) / 2.
    center_lon = (rad_image.lon_range[0] + rad_image.lon_range[1]) / 2.
    rad_image.grid_x, rad_image.grid_y = _latlon_to_xy(
        rad_image.grid_lat, rad_image.grid_lon, center_lat, center_lon)
    rad_image.times = times
    return rad_image

def _preprocess(rad_file, lat_range, lon_range):
    radar = pyart.io.read(rad_file)
    disp = pyart.graph.RadarMapDisplay(radar)
    fig, ax = plt.subplots(1, 1, figsize=(2.56, 2.56),
            subplot_kw=dict(projection=ccrs.PlateCarree(), frameon=False))

    disp.plot_ppi_map('reflectivity', sweep=0, min_lon=lon_range[0],
            ax=ax, max_lon=lon_range[1], min_lat=lat_range[0], max_lat=lat_range[1],
            embellish=False, vmin=-20, vmax=60, cmap='HomeyerRainbow',
            add_grid_lines=False, colorbar_flag=False, title_flag=False)
    ax.set_axis_off()
    fig.tight_layout(pad=0, w_pad=0, h_pad=0)
    with tempfile.NamedTemporaryFile(mode='w+b' ) as temp_file:
        fig.savefig(temp_file, dpi=100)
        plt.close(fig)
        # Transform image
        image = decode_image(temp_file.name)
    # Transform image
    image = image[:3, :, :].float()
    transform = transforms.Compose([
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    image = torch.stack([transform(image)])
    rad_time = np.datetime64(radar.time["units"].split()[2])
    del radar
    return image, rad_time

def _latlon_to_xy(lat, lon, lat0=0, lon0=0):
    R = 6371000  # Earth's radius in meters
    x = np.radians(lon - lon0) * R * np.cos(np.radians(lat0))
    y = np.radians(lat - lat0) * R
    return x, y
