import os
import matplotlib.pyplot as plt
from datetime import datetime
import cftime
import time

import numpy as np
import xarray as xr
import pandas as pd
import matplotlib.dates as mdates

class Plotter:
    def __init__(self, data_array, latitudes, longitudes, datasets, output_dir):
        self.data_array = data_array
        self.latitudes = latitudes
        self.longitudes = longitudes
        self.datasets = datasets
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)

    def convert_cftime_to_datetime(self, times):
        return [
            datetime(t.year, t.month, t.day, t.hour, t.minute, t.second)
            if isinstance(t, cftime.DatetimeNoLeap) else t
            for t in times
        ]

    def plot_with_markers(self, var, time_index, index_pairs):
        # Existing implementation of plot_with_markers
        """Plot data with markers at specified index pairs."""
        # Select the desired time slice
        data = self.data_array.isel(time=time_index)

        # Loop through each pair of indices and plot the markers
        print("self.data_array",self.data_array)
        print("index_pairs ", index_pairs)

        print("(datetime.now()..........",datetime.now().strftime("%Y-%m-%d_%HH:%MM:%SS"))
        # Extract the list of variables to plot from the datasets
        list_multivar = list(self.datasets.keys())
        print("list_multivar...",list_multivar)

        for var_idx, var_name in enumerate(list_multivar):
                dataset = self.datasets.get(var_name)
                print("var_name.......", var_name)
                print("var_idx...........", var_idx)

                # Plot the data
                if dataset is not None:

                    image = self.datasets.get(var_name).isel(time=time_index)


                    print("image ",image)
                    print("image.dims",image.dims)
                    print("len image.dims",len(image.dims))
                    # Reduce dimensionality if necessary
                    if 'lev' in image.dims:
                        image = image.isel(lev=6)  # Adjust level index as needed

                    if 'vertices' in image.dims:
                        image = image.mean(dim='vertices')

                    if 'bnds' in image.dims:
                        image = image.mean(dim='bnds')

                    # Check if 'image' is a Dataset
                    if isinstance(image, xr.Dataset):
                        # Select the variable within the dataset
                        image = image[var_name.lower()]  # Make sure this matches the variable name in the Dataset

                    print("AFTER Reduction image.dims",image.dims)
                    print("AFTER Reduction len image.dims",len(image.dims))


                    # Check that the data is now 2-dimensional
                    if len(image.dims) == 2:
                        plot = image.plot.pcolormesh()
                    else:
                        raise ValueError("Data is not 2-dimensional after reduction.")





                    # Extract the axis from the plot
                    ax = plot.axes
                    # Loop through each pair of indices and plot the markers
                    k=0
                    for i_index, j_index in index_pairs:

                        k += 1



                        value = self.data_array.isel(time=time_index, j=j_index, i=i_index).values
                        lat = self.latitudes.isel(i=i_index, j=j_index).values
                        lon = self.longitudes.isel(i=i_index, j=j_index).values
                        print("i_index ",i_index)
                        print("j_index ",j_index)
                        print("lat ",lat)
                        print("lon ",lon)
                        print("type value ",type(value))
                        print("value ",type(self.data_array.isel(time=time_index, j=j_index, i=i_index).values))
                        #print(f'i:{i_index}, j:{j_index}, lat:{lat:.1f}, lon:{lon:.1f}, value:{value:.2e}')
                        # Add marker to the plot
                        ax.plot(i_index, j_index, marker='x', color='red', markersize=10)
                        ax.set_title(f'i:{i_index}, j:{j_index}, lat:{lat:.1f}, lon:{lon:.1f}')

                        # Save the plot as an image
                        output_path = f'{self.output_dir}/{var_name}_images_{k}_location_{i_index}_{j_index}.png'

                        os.makedirs(self.output_dir, exist_ok=True)
                        plt.savefig(output_path)

                        # Optionally print details to console
                        #print(f'Index: ({i_index}, {j_index}), Value: {value:.2e}, Latitude: {lat:.1f}, Longitude: {lon:.1f}')

                        # Show the final plot
                        #plt.show()
                    plt.close()

                else:
                        raise ValueError("Data is Empty")


    def plot_series(self, index_pairs):
        """Plot data series as subplots."""
        start_time_plot1 = time.time()
        #print("start_time_plot1........", start_time_plot1)
        print("(datetime.now()..........",datetime.now().strftime("%Y-%m-%d_%HH:%MM:%SS"))
        # Extract the list of variables to plot from the datasets
        list_multivar = list(self.datasets.keys())
        print("list_multivar...",list_multivar)


        # Loop over each index pair
        for idx, (i_index, j_index) in enumerate(index_pairs):

            print(i_index, j_index)
            fig, axs = plt.subplots(
                len(list_multivar), 1, figsize=(15, 3 * len(list_multivar)),
                constrained_layout=True
            )

            # If only one subplot, wrap axs in a list to make it iterable
            if len(list_multivar) == 1:
                axs = [axs]

            for var_idx, var_name in enumerate(list_multivar):
                dataset = self.datasets.get(var_name)
                print("var_name.......", var_name)
                print("var_idx...........", var_idx)
                if dataset is not None:
                    # Check the level selection based on the variable name
                    if var_name in ['INTPP', 'MLOTST', 'ZMESOOS', 'ZMICROOS', 'RSNTDS']:
                        timeseries = self.datasets.get(var_name).isel(i=i_index, j=j_index)
                        lat = dataset.latitude.isel(i=i_index, j=j_index).values
                        lon = dataset.longitude.isel(i=i_index, j=j_index).values
                    else:
                        timeseries = self.datasets.get(var_name).isel(i=i_index, j=j_index, lev=6)
                        lat = dataset.latitude.isel(i=i_index, j=j_index).values
                        lon = dataset.longitude.isel(i=i_index, j=j_index).values
                    # Debug check timeseries
                    #print("timeseries.........",timeseries)
                    #print("timeseries.latitude.values............... ",timeseries.latitude.values)
                    # Debug check timeseries
                    #print("timeseries[var_name.lower()].values........... ",timeseries[var_name.lower()].values)
                    #print("np.shape(timeseries[var_name.lower()].values)...... ", np.shape(timeseries[var_name.lower()].values))
                    # Convert 'time' to standard datetime if needed
                    times = self.convert_cftime_to_datetime(timeseries['time'].values)
                    # Debug check times
                    #print("times......",times)
                    #print("times......",np.shape(times))



                    # Plot the timeseries directly
                    axs[var_idx].plot(times, timeseries[var_name.lower()].values)
                    axs[var_idx].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))

                    axs[var_idx].set_title(
                        f'{var_name} for i={i_index}, j={j_index}, lat={lat:.0f}, lon={lon:.0f}'
                    )
                    axs[var_idx].set_xlabel('Time')
                    axs[var_idx].set_ylabel('Value')
                    axs[var_idx].grid(True)
                else:
                    axs[var_idx].text(
                        0.5, 0.5, f'{var_name} dataset not found',
                        horizontalalignment='center', verticalalignment='center'
                    )
                    axs[var_idx].set_title(
                        f'{var_name} for i={i_index}, j={j_index}'
                    )
                    axs[var_idx].set_xlabel('Time')
                    axs[var_idx].set_ylabel('Value')
                    axs[var_idx].grid(True)

            plt.suptitle(f'Timeseries for Location (i={i_index}, j={j_index})', fontsize=16)
            plt.savefig(f'{self.output_dir}/timeseries_location_{i_index}_{j_index}.png')
            plt.close(fig)

        print("All plots have been saved.")
        end_time_plot1 = time.time()
        time_used_plot1 = (end_time_plot1 - start_time_plot1)
        print("Time used for entire PLOT..../zzzz", f'{time_used_plot1:.4f}', "Seconds")
        print("Time used for entire PLOT..../zzzz", f'{time_used_plot1/60:.4f}', "Minutes")
        print("Time used for entire PLOT..../zzzz", f'{time_used_plot1/3600:.4f}', "Hours")
