import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
import logging
import argparse
from datetime import datetime
import zarr
import shutil
import json
import sys  # Added to get the size of the objects
import json

def plot_data(data, title, filename):
    plt.figure(figsize=(10, 6))
    plt.imshow(data, cmap='viridis')
    plt.colorbar(label='Value')
    plt.title(title)
    plt.savefig(filename)
    plt.close()

def get_size(obj, name):
    size = sys.getsizeof(obj)
    logging.info(f"Size of {name}: {size} bytes")
    return size

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Process land use data and generate yearly datasets.")
parser.add_argument("--start_year", type=int, required=True, help="Start year for the data range")
parser.add_argument("--end_year", type=int, required=True, help="End year for the data range")
parser.add_argument("--region", type=str, required=True, help="Region for the data slice (ES, EU, UK, AU, GLOBE)")
parser.add_argument("--include_luh2", action='store_true', help="Include LUH2 dataset in the processed data")
parser.add_argument("--include_hilda", action='store_true', help="Include HILDA+ dataset in the processed data")
parser.add_argument("--input_data_override", type=str, help="Path to override the default input data directory")
parser.add_argument("--default_data_path", type=str, required=True, help="Default input data directory")
parser.add_argument("--save_data_path", type=str, required=True, help="Directory to save processed data")
parser.add_argument("--inspection", action='store_true', help="Enable inspection mode with additional plotting")
parser.add_argument("--overwrite", action='store_true', help="Overwrite existing processed data if it exists")
parser.add_argument("--dtype_ground_truth", type=str, default="int16", help="Data type for ground truth data (e.g., int16, float32)")
parser.add_argument("--dtype_luh2", type=str, default="float32", help="Data type for LUH2 data (e.g., int16, float32)")
parser.add_argument("--dtype_hilda", type=str, default="int16", help="Data type for HILDA+ data (e.g., int16, float32)")
parser.add_argument("--regions_file", type=str, required=True, help="Path to the regions JSON file")

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
region = args.region
include_luh2 = args.include_luh2
include_hilda = args.include_hilda
input_data_override = args.input_data_override
inspection_mode = args.inspection
overwrite_existing = args.overwrite
dtype_ground_truth = args.dtype_ground_truth
dtype_luh2 = args.dtype_luh2
dtype_hilda = args.dtype_hilda
default_data_path = args.default_data_path
save_data_path = args.save_data_path
regions_file_path = args.regions_file

# Configure logging
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

# Create a directory for logs
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)

log_file_name = f"preprocessing_log_{start_year}_{end_year}_{region}_{timestamp}.log"
log_file_path = os.path.join(logs_path, log_file_name)

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[
    logging.FileHandler(log_file_path),
    logging.StreamHandler()
])

logging.info("Preprocessing started.")
logging.info(f"Arguments received: start_year={start_year}, end_year={end_year}, region={region}, include_luh2={include_luh2}, include_hilda={include_hilda}, input_data_override={input_data_override}, inspection_mode={inspection_mode}, overwrite_existing={overwrite_existing}, dtype_ground_truth={dtype_ground_truth}, dtype_luh2={dtype_luh2}, dtype_hilda={dtype_hilda}, default_data_path={default_data_path}, save_data_path={save_data_path}, regions_file_path={regions_file_path}")

# Define input data path
input_data_path = input_data_override if input_data_override else default_data_path
logging.info(f"Using input data path: {input_data_path}")

# Check if the regions file exists
if not os.path.exists(regions_file_path):
    logging.error(f"Regions JSON file not found at: {regions_file_path}")
    sys.exit(1)

# Load the region information from the JSON file
with open(regions_file_path, 'r') as regions_file:
    regions_data = json.load(regions_file)

# Check if the specified region exists in the regions data
if region not in regions_data:
    logging.error(f"Region {region} not found in regions JSON file.")
    sys.exit(1)

# Extract latitude and longitude slices for the given region
lon_slice = slice(*regions_data[region]["lon_slice"])
lat_slice = slice(*regions_data[region]["lat_slice"])

logging.info(f"Set slicing for region {region}: lon_slice={lon_slice}, lat_slice={lat_slice}")

# Create a main directory for plots if it doesn't exist
plots_main_path = os.path.join(save_data_path, 'plots')
os.makedirs(plots_main_path, exist_ok=True)
logging.info(f"Plots main directory set at: {plots_main_path}")

# Create a region-specific subdirectory for plots
plots_path = os.path.join(plots_main_path, f'{region}_{start_year}_{end_year}')
os.makedirs(plots_path, exist_ok=True)
logging.info(f"Plots subdirectory set at: {plots_path}")

# Initialize a dictionary to store metadata
metadata = {
    "start_year": start_year,
    "end_year": end_year,
    "region": region,
    "include_luh2": include_luh2,
    "include_hilda": include_hilda,
    "data_path": input_data_path,
    "processing_date": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
    "dtype_ground_truth": dtype_ground_truth,
    "dtype_luh2": dtype_luh2,
    "dtype_hilda": dtype_hilda
}

# Loop through each year and process data
for year in range(start_year, end_year + 1):
    logging.info(f"Processing year: {year}")

    # Load the ground truth dataset
    ground_truth_path = os.path.join(input_data_path, f"lccsclass/lccsclass_{year}_0.0089deg.nc")
    logging.info(f"Loading ground truth data from: {ground_truth_path}")

    # Paths to the LUH2 datasets
    # TODO : The following dataset should be removed : "secmb", "secma"
    luh2_data_paths = [os.path.join(input_data_path, f"luh2class/{dataset}_{year}_0.0089deg.nc") for dataset in [
        #"primf", "primn", "secdf", "secdn", "urban", "c3ann", "c4ann", "c3per", "c4per", "c3nfx", "pastr", "range", "secmb", "secma"
        "primf", "primn", "secdf", "secdn", "urban", "c3ann", "c4ann", "c3per", "c4per", "c3nfx", "pastr", "range"
    ]]
    logging.info(f"LUH2 data paths: {luh2_data_paths}")

    # Path to the HILDA+ dataset
    hilda_dataset_path = os.path.join(input_data_path, f"hilda/hildaplus_GLOB-1-0-f_states_{year}_0.0089deg.nc")
    logging.info(f"HILDA+ dataset path: {hilda_dataset_path}")

    # Open the ground truth dataset with slicing
    ground_truth_ds = xr.open_dataset(ground_truth_path).sel(lon=lon_slice, lat=lat_slice)
    logging.info(f"Ground truth dataset loaded and sliced with dimensions: {ground_truth_ds.dims}")

    # Convert the ground truth data to the specified dtype
    ground_truth_ds['lccs_class'] = ground_truth_ds['lccs_class'].astype(dtype_ground_truth)
    logging.info(f"Ground truth data converted to dtype: {dtype_ground_truth}")
    get_size(ground_truth_ds['lccs_class'].values, 'Ground Truth Data')

    # Get actual dimensions for reshaping using sizes
    lat_dim = ground_truth_ds.sizes['lat']
    lon_dim = ground_truth_ds.sizes['lon']
    logging.info(f"lat_dim={lat_dim}, lon_dim={lon_dim}")

    # Save these dimensions to a text file or a numpy file
    dims_path = os.path.join(save_data_path, f'dims_{start_year}_{end_year}_{region}.npy')
    np.save(dims_path, np.array([lat_dim, lon_dim]))
    logging.info(f"Saved dimensions to: {dims_path}")

    # Update metadata with dimensions
    metadata["lat_dim"] = int(lat_dim)
    metadata["lon_dim"] = int(lon_dim)

    # Flip the ground truth data (both upside-down and left-to-right twice)
    ground_truth_ds['lccs_class'] = ground_truth_ds['lccs_class'].transpose('lat', 'lon', 'time')
    ground_truth_ds['lccs_class'].values = np.flipud(ground_truth_ds['lccs_class'].values)
    logging.info("Ground truth data flipped and transposed.")

    # Select the first time step
    lccs_class_first_time = ground_truth_ds['lccs_class'].isel(time=0)
    logging.info(f"Selected the first time step for LCCS class data with shape: {lccs_class_first_time.shape}")

    # Plot the ground truth data only if inspection mode is enabled
    if inspection_mode:
        logging.info("Plotting ground truth data for inspection mode.")
        plot_data(
            lccs_class_first_time.values,
            f"LCCS Class - {region} {year}",
            os.path.join(plots_path, f"lccs_class_{region}_{year}_before.png")
        )

    # Initialize list for features
    input_features = []

    if include_luh2:
        # Open the LUH2 datasets and combine them
        luh2_ds = xr.open_mfdataset(luh2_data_paths, combine='by_coords').sel(lon=lon_slice, lat=lat_slice)
        logging.info(f"LUH2 datasets loaded and sliced with dimensions: {luh2_ds.dims}")

        # Convert LUH2 data to the specified dtype
        for var in luh2_ds.data_vars:
            luh2_ds[var] = luh2_ds[var].astype(dtype_luh2)
        logging.info(f"LUH2 data converted to dtype: {dtype_luh2}")
        get_size(luh2_ds, 'LUH2 Data')

        # Flip the LUH2 data (both upside-down and left-to-right twice)
        for var in luh2_ds.data_vars:
            luh2_ds[var].values = np.flipud(luh2_ds[var].values)
        logging.info("LUH2 data flipped.")

        # Plot each LUH2 dataset variable only if inspection mode is enabled
        if inspection_mode:
            for var in luh2_ds.data_vars:
                logging.info(f"Plotting LUH2 data for variable: {var}")
                plot_data(
                    luh2_ds[var].values,
                    f"LUH2 {var} - {region} {year}",
                    os.path.join(plots_path, f"luh2_{var}_{region}_{year}_before.png")
                )

        # Append LUH2 features
        input_features.extend([luh2_ds[var].values.flatten() for var in luh2_ds.data_vars])
        logging.info(f"LUH2 features appended. Number of features: {len(input_features)}")

    if include_hilda:
        # Open the HILDA+ dataset
        hilda_ds = xr.open_dataset(hilda_dataset_path, decode_times=False).sel(lon=lon_slice, lat=lat_slice)
        logging.info(f"HILDA+ dataset loaded with dimensions: {hilda_ds.dims}")

        # Convert HILDA+ data to the specified dtype
        hilda_ds['LULC_states'] = hilda_ds['LULC_states'].astype(dtype_hilda)
        logging.info(f"HILDA+ data converted to dtype: {dtype_hilda}")
        get_size(hilda_ds['LULC_states'].values, 'HILDA+ Data')

        # Remove classification information from HILDA+ and only use water data (classification code 0 for ocean)
        # hilda_ds['LULC_states'] = xr.where((hilda_ds['LULC_states'] == 77) | (hilda_ds['LULC_states'] == 0), 0, 1) this was removed as in-land water is excluded
        hilda_ds['LULC_states'] = xr.where(hilda_ds['LULC_states'] == 0, 0, 1)
        logging.info("Non-water classifications set to 1; water classifications set to 0.")

        # # Normalize HILDA+ data to range [0, 1] using the maximum value found in the dataset
        # max_value_hilda = hilda_ds['LULC_states'].max().values  # Compute maximum value dynamically
        # hilda_ds['LULC_states'] = hilda_ds['LULC_states'] / max_value_hilda  # Normalize dynamically
        # logging.info(f"HILDA+ data normalized to range [0, 1] with max value {max_value_hilda}.")

        # Flip the HILDA+ data (both upside-down and left-to-right twice)
        hilda_ds_flipped = np.flipud(hilda_ds['LULC_states'].values.squeeze())
        logging.info("HILDA+ data flipped.")

        # Create a DataArray with explicit dimension names
        hilda_ds['LULC_states'] = xr.DataArray(hilda_ds_flipped, dims=['lat', 'lon'], coords={'lat': hilda_ds['lat'], 'lon': hilda_ds['lon']})

        # Plot the HILDA+ data only if inspection mode is enabled
        if inspection_mode:
            logging.info("Plotting HILDA+ data for inspection mode.")
            plot_data(
                hilda_ds['LULC_states'].values,
                f"HILDA+ LULC States - {region} {year}",
                os.path.join(plots_path, f"hilda_lulc_states_{region}_{year}_before.png")
            )

        # Append normalized HILDA+ features
        input_features.append(hilda_ds['LULC_states'].values.flatten())
        logging.info("HILDA+ features appended.")

    # Ensure all features have the same length before stacking
    if input_features:  # Check if input_features is not empty
        min_length = min([len(f) for f in input_features] + [len(lccs_class_first_time.values.flatten())])
        input_features = [f[:min_length] for f in input_features]
        lccs_class = lccs_class_first_time.values.flatten()[:min_length]
        logging.info(f"Data truncated to minimum length: {min_length}")

        # Save each year's processed data separately
        lccs_class_year_path = os.path.join(save_data_path, f'lccs_class_{year}_{region}.zarr')
        features_year_path = os.path.join(save_data_path, f'features_{year}_{region}.zarr')

        zarr.save(lccs_class_year_path, lccs_class)
        zarr.save(features_year_path, np.vstack(input_features).T)
        logging.info(f"Data for year {year} saved at: {lccs_class_year_path} and {features_year_path}")
    else:
        logging.warning(f"No input features were found for year {year}. Skipping this year.")

# Save metadata including dimensions, included datasets, and other relevant details
metadata_path = os.path.join(save_data_path, f'metadata_{start_year}_{end_year}_{region}.json')
with open(metadata_path, 'w') as f:
    json.dump(metadata, f)
logging.info(f"Metadata saved to: {metadata_path}")

logging.info("Preprocessing completed successfully.")
