from PIL import Image
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import MinMaxScaler
import joblib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import gradio as gr
import tensorflow as tf
import data_processing
import prediction
from huggingface_hub import snapshot_download

# Load pre-trained models and data
path = snapshot_download("abatejemal/3_Models")
ohe_loaded = joblib.load(f'{path}/transform_ohe.pkl')
df = pd.read_csv('bbox_and_commons.csv')
districts = df['district'].tolist()

# Load district-specific data and models
def load_district_data(district_selected):
    dataset_paths = f"https://huggingface.co/datasets/abatejemal/2_Data/resolve/main/WeatherData/{district_selected}.csv"
    dataset_path = dataset_paths
    model_paths = f"3_Models/weather_models/{district_selected}_lstm_model.h5"
    scaler_paths = f"3_Models/weather_models/{district_selected}_scaler.pkl"
    
    data = pd.read_csv(dataset_path)
    data['date'] = pd.to_datetime(data['date'])
    data.set_index('date', inplace=True)
    numeric_columns = ['GWETPROF', 'GWETTOP', 'GWETROOT', 'CLOUD_AMT', 'TS', 'PS', 'RH2M', 'QV2M', 'PRECTOTCORR', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE', 'WS2M']
    data = data[numeric_columns].dropna()
    data = data_processing.fill_outliers_with_median(data)
    
    return data

# Prediction function
def predict_crop_yield(district_selected, area, selected_season):
    data = load_district_data(district_selected)
    
    # Scale data and predict
    time_steps = 365
    predictions = prediction.predict_next_30_days(district_selected, data, time_steps, days=90)
    predicted_df = pd.DataFrame(predictions, columns=['GWETPROF', 'GWETTOP', 'GWETROOT', 'CLOUD_AMT', 'TS', 'PS', 'RH2M', 'QV2M', 'PRECTOTCORR', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE', 'WS2M'])
    
    # Calculate mean of predictions
    mean_predicted_df = predicted_df.mean(numeric_only=True)
    mean_row = pd.DataFrame([mean_predicted_df.tolist()], columns=predicted_df.columns)
    
    filtered_df = df[df['district'] == district_selected]
    filtered_df = filtered_df.reset_index(drop=True)
    
    mean_row['elevation'] = filtered_df['elevation']
    mean_row['slope'] = filtered_df['slope']
    mean_row['soc'] = filtered_df['soc']
    mean_row['soilph'] = filtered_df['soilph']
    mean_row['area(sq.m)'] = area
    mean_row['season'] = selected_season
    
    important_columns = ['season','crop', 'area(sq.m)', 'GWETPROF', 'GWETTOP', 'GWETROOT', 'CLOUD_AMT', 'TS', 'PS', 'RH2M', 'QV2M', 'PRECTOTCORR', 'T2M_MAX', 'T2M_MIN', 'T2M_RANGE', 'WS2M', 'elevation', 'slope', 'soc', 'soilph']
    ch = pd.DataFrame()
    crop_categories = ohe_loaded.categories_[0] 
    
    for crop in crop_categories:
        d = mean_row
        d['crop'] = crop
        ch = pd.concat([ch, d])
    
    final = ch[important_columns]
    final = final.reset_index(drop=True)
    encoded_final = ohe_loaded.transform(final[['crop', 'season']])
    encoded_final = pd.DataFrame(encoded_final.toarray(), 
              columns=[f"{val}" for cat, vals in zip(ohe_loaded.feature_names_in_, ohe_loaded.categories_) for val in vals])
    
    final_df = pd.concat([final[['area(sq.m)', 'GWETPROF', 'GWETTOP', 'GWETROOT', 'CLOUD_AMT', 
           'TS', 'PS', 'RH2M', 'QV2M', 'PRECTOTCORR', 'T2M_MAX', 
           'T2M_MIN', 'T2M_RANGE', 'WS2M', 'elevation', 'slope', 'soc', 'soilph']], encoded_final], axis=1)
    
    predicted_df = prediction.predict_crop_yield1(final, encoded_final, ohe_loaded)
    season = predicted_df['season'].tolist()
    tstr = f'Predicted Production in District: {district_selected}   Season: {season[0]}'
    
    fig = plt.figure(figsize=(10, 6))
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_title(tstr, fontsize=15)
    ax.set_ylabel('Production', fontsize=14)
    ax.set_xlabel('Crop', fontsize=13)
    ax.bar(predicted_df['crop'][:8], predicted_df['Predicted'][:8])
    
    return fig, predicted_df

# Gradio Interface
def gradio_interface(district_selected, area, selected_season):
    with gr.Blocks() as demo:
        district_dropdown = gr.Dropdown(label="Select District", choices=districts)
        season_dropdown = gr.Dropdown(label="Select Season", choices=["Meher", "Belg"])
        area_input = gr.Number(label="Enter Area (sq.m)", min_value=10, max_value=4000)
        submit_button = gr.Button("Predict")
        
        output_plot = gr.Plot(label="Prediction Plot")
        output_df = gr.DataFrame(label="Prediction Results")
        
        submit_button.click(predict_crop_yield, inputs=[district_dropdown, area_input, season_dropdown], outputs=[output_plot, output_df])
        
    return demo.launch()

# Launch the Gradio app
gradio_interface(districts[0], 1000, "Meher")