import xarray as xr
import numpy as np
import os
import logging
import argparse
from datetime import datetime
import zarr
import json
import sys
import matplotlib.pyplot as plt

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="Inference preprocessing script for generating features.")
parser.add_argument("--start_year", type=int, required=True, help="Start year for inference")
parser.add_argument("--end_year", type=int, required=True, help="End year for inference")
parser.add_argument("--train_start_year", type=int, required=True, help="Start year of the training data")
parser.add_argument("--train_end_year", type=int, required=True, help="End year of the training data")
parser.add_argument("--region", type=str, required=True, help="Region code (e.g., 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("--save_data_path", type=str, required=True, help="Path to save processed data")
parser.add_argument("--default_data_path", type=str, required=True, help="Path for default input data")
parser.add_argument("--inspection", action='store_true', help="Enable inspection mode with additional plotting")
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")
parser.add_argument("--read_db", action="store_true", help="Read model details from the database")
parser.add_argument("--db_path", type=str, required=False, help="Path to the SQLite database")

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
save_data_path = args.save_data_path
default_data_path = args.default_data_path
inspection_mode = args.inspection
dtype_luh2 = args.dtype_luh2
dtype_hilda = args.dtype_hilda
train_start_year = args.train_start_year
train_end_year = args.train_end_year
regions_file_path = args.regions_file
read_database = args.read_db
databashe_path = args.db_path

# Configure logging
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)
log_file_name = f"inference_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("Inference 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}, save_data_path={save_data_path}, inspection_mode={inspection_mode}, dtype_luh2={dtype_luh2}, dtype_hilda={dtype_hilda}, train_start_year={train_start_year}, train_end_year={train_end_year}, regions_file_path={regions_file_path}")

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

# Load region slicing from a JSON file
if not os.path.exists(regions_file_path):
    logging.error(f"Regions JSON file not found at: {regions_file_path}")
    sys.exit(1)

with open(regions_file_path, 'r') as regions_file:
    regions_data = json.load(regions_file)

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}")

# 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_luh2": dtype_luh2,
    "dtype_hilda": dtype_hilda
}

# Create the plots directory inside save_data_path
if inspection_mode:
    plots_path = os.path.join(save_data_path, "plots", f"{region}_{start_year}_{end_year}")
    os.makedirs(plots_path, exist_ok=True)

# Load dimensions from the file based on training years
dims_path = os.path.join(save_data_path, f'dims_{train_start_year}_{train_end_year}_{region}.npy')
if os.path.exists(dims_path):
    lat_dim, lon_dim = np.load(dims_path)
    logging.info(f"Loaded dimensions from existing file: lat_dim={lat_dim}, lon_dim={lon_dim}")
    metadata["lat_dim"] = int(lat_dim)
    metadata["lon_dim"] = int(lon_dim)
else:
    logging.error(f"Dimensions file not found at: {dims_path}")
    sys.exit(1)

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

    # Paths to the LUH2 datasets
    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"
    ]]
    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}")

    # 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 (upside-down)
        for var in luh2_ds.data_vars:
            luh2_ds[var].values = np.flipud(luh2_ds[var].values)
        logging.info("LUH2 data flipped.")

        # Plot LUH2 data if inspection mode is enabled
        if inspection_mode:
            for var in luh2_ds.data_vars:
                plot_data(
                    luh2_ds[var].values,
                    f"LUH2 {var} - {region} {year}",
                    os.path.join(plots_path, f"luh2_{var}_{region}_{year}.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}")

        # Keep a copy of the original HILDA+ LULC_states
        hilda_ds['HILDA_original_lulc_states'] = hilda_ds['LULC_states'].copy()
        logging.info("A copy of the original HILDA+ LULC_states has been kept as 'HILDA_original_lulc_states'.")

        # 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/open sea and 77 for inland water)
        hilda_ds['LULC_states'] = xr.where((hilda_ds['LULC_states'] == 77) | (hilda_ds['LULC_states'] == 0), 0, 1)
        logging.info("Non-water classifications set to 1; water classifications set to 0.")

        # Flip the processed HILDA+ data (upside-down)
        hilda_ds_flipped = np.flipud(hilda_ds['LULC_states'].values.squeeze())
        logging.info("Processed HILDA+ data flipped.")

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

        # Plot the processed HILDA+ data only if inspection mode is enabled
        if 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}_processed.png")
            )

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

        # Process the original HILDA+ data (HILDA_original_lulc_states)
        hilda_original_flipped = np.flipud(hilda_ds['HILDA_original_lulc_states'].values.squeeze())
        logging.info("Original HILDA+ data flipped.")

        # Create a DataArray with explicit dimension names for the original flipped data
        hilda_ds['HILDA_original_lulc_states'] = xr.DataArray(hilda_original_flipped, dims=['lat', 'lon'], coords={'lat': hilda_ds['lat'], 'lon': hilda_ds['lon']})

        # Plot the original HILDA+ data only if inspection mode is enabled
        if inspection_mode:
            plot_data(
                hilda_ds['HILDA_original_lulc_states'].values,
                f"Original HILDA+ LULC States - {region} {year}",
                os.path.join(plots_path, f"hilda_original_lulc_states_{region}_{year}.png")
            )

        # Save original HILDA+ features separately
        original_hilda_features_path = os.path.join(save_data_path, f'original_hilda_features_{year}_{region}_inference.zarr')

        # Save the original HILDA+ features using Zarr
        zarr.save(original_hilda_features_path, hilda_ds['HILDA_original_lulc_states'].values)
        logging.info(f"Original HILDA+ features for year {year} saved separately at: {original_hilda_features_path}")


    # 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])
        input_features = [f[:min_length] for f in input_features]
        logging.info(f"Data truncated to minimum length: {min_length}")

        # Save features for the year
        features_year_path = os.path.join(save_data_path, f'features_{year}_{region}_inference.zarr')

        # Save features using Zarr
        zarr.save(features_year_path, np.vstack(input_features).T)
        logging.info(f"Features for year {year} saved at: {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}_inference.json')
with open(metadata_path, 'w') as f:
    json.dump(metadata, f)
logging.info(f"Metadata saved to: {metadata_path}")

logging.info("Inference preprocessing completed successfully.")
