import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.svm import SVR
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.preprocessing import LabelEncoder
import gradio as gr
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Input
from tensorflow.keras.optimizers import Adam
from PIL import Image
import rasterio
import matplotlib.pyplot as plt
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.models import Model
import cv2
import joblib

# Load crop data
def load_data():
    url = 'https://raw.githubusercontent.com/NarutoOp/Crop-Recommendation/master/cropdata.csv'
    data = pd.read_csv(url)
    return data

data = load_data()

label_encoders = {}
for column in ['STATE', 'CROP']:
    le = LabelEncoder()
    data[column] = le.fit_transform(data[column])
    label_encoders[column] = le

X = data[['YEAR', 'STATE', 'CROP', 'YEILD']]  # Feature columns
y = data['PROFIT']  # Target column

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

models = {
    'Linear Regression': LinearRegression(),
    'SVR': SVR(),
    'Random Forest': RandomForestRegressor(),
    'Gradient Boosting': GradientBoostingRegressor()
}

for name, model in models.items():
    model.fit(X_train, y_train)

def predict_traditional(model_name, year, state, crop, yield_):
    if model_name in models:
        model = models[model_name]
        state_encoded = label_encoders['STATE'].transform([state])[0]
        crop_encoded = label_encoders['CROP'].transform([crop])[0]
        prediction = model.predict([[year, state_encoded, crop_encoded, yield_]])[0]
        return prediction
    else:
        return "Model not found"

# Train RandomForestRegressor model for deep learning model
def train_random_forest_model():
    def process_tiff(file_path):
        with rasterio.open(file_path) as src:
            tiff_data = src.read()
        B2_image = tiff_data[1, :, :]  # Assuming B2 is the second band
        target_size = (50, 50)
        B2_resized = cv2.resize(B2_image, target_size, interpolation=cv2.INTER_NEAREST)
        return B2_resized.reshape(-1, 1)

    data_dir = 'Data'
    X_list = []
    y_list = []

    for root, dirs, files in os.walk(data_dir):
        for file in files:
            if file.endswith('.tiff'):
                file_path = os.path.join(root, file)
                X_list.append(process_tiff(file_path))
                y_list.append(np.random.rand(2500))  # Replace with actual target data

    X = np.vstack(X_list)
    y = np.hstack(y_list)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    return model

rf_model = train_random_forest_model()

def predict_random_forest(file):
    if file is not None:
        def process_tiff(file_path):
            with rasterio.open(file_path) as src:
                tiff_data = src.read()
            B2_image = tiff_data[1, :, :]
            target_size = (50, 50)
            B2_resized = cv2.resize(B2_image, target_size, interpolation=cv2.INTER_NEAREST)
            return B2_resized.reshape(-1, 1)

        tiff_processed = process_tiff(file.name)
        prediction = rf_model.predict(tiff_processed)
        prediction_reshaped = prediction.reshape((50, 50))

        plt.figure(figsize=(10, 10))
        plt.imshow(prediction_reshaped, cmap='viridis')
        plt.colorbar()
        plt.title('Yield Prediction for Single TIFF File')
        plt.savefig('/tmp/rf_prediction_overlay.png')

        return '/tmp/rf_prediction_overlay.png'
    else:
        return "No file uploaded"

# Load deep learning models
def load_deep_learning_model(model_name):
    base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(128, 128, 3))
    base_model.trainable = False

    inputs = Input(shape=(128, 128, 3))
    x = base_model(inputs, training=False)
    x = GlobalAveragePooling2D()(x)
    outputs = Dense(1, activation='linear')(x)

    model = Model(inputs, outputs)
    model.compile(optimizer=Adam(), loss='mean_squared_error', metrics=['mae'])

    return model

deep_learning_models = {
    'ResNet50': load_deep_learning_model('ResNet50'),
    # Add other models here if needed
}

def predict_deep_learning(model_name, file):
    if model_name in deep_learning_models:
        if file is not None:
            with rasterio.open(file.name) as src:
                img_data = src.read(1)

            patch_size = 128
            n_patches_x = img_data.shape[1] // patch_size
            n_patches_y = img_data.shape[0] // patch_size

            patches = []
            for i in range(n_patches_y):
                for j in range(n_patches_x):
                    patch = img_data[i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size]
                    patches.append(patch)
            patches = np.array(patches)

            preprocessed_patches = []
            for patch in patches:
                img = Image.fromarray(patch)
                img = img.convert('RGB')
                img = img.resize((128, 128))
                img_array = np.array(img) / 255.0
                preprocessed_patches.append(img_array)
            preprocessed_patches = np.array(preprocessed_patches)

            model = deep_learning_models[model_name]
            predictions = model.predict(preprocessed_patches)
            predictions = predictions.reshape((n_patches_y, n_patches_x))

            # Set a threshold to highlight areas with higher predicted yields
            threshold = np.percentile(predictions, 90)  # Adjust the percentile as needed

            # Create an overlay image to visualize predictions
            overlay = np.zeros_like(img_data, dtype=np.float32)
            for i in range(n_patches_y):
                for j in range(n_patches_x):
                    if predictions[i, j] > threshold:
                        overlay[i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size] = predictions[i, j]

            # Plot the overlay on the original image
            plt.figure(figsize=(10, 10))
            plt.imshow(img_data, cmap='gray', alpha=0.5)
            plt.imshow(overlay, cmap='jet', alpha=0.5)
            plt.title('Crop Yield Prediction Overlay')
            plt.colorbar()

            # Save the plot to a file
            plt.savefig('/tmp/dl_prediction_overlay.png')

            return '/tmp/dl_prediction_overlay.png'
        else:
            return "No file uploaded"
    else:
        return "Model not found"

inputs_traditional = [
    gr.Dropdown(choices=list(models.keys()), label='Model'),
    gr.Number(label='Year'),
    gr.Textbox(label='State'),
    gr.Textbox(label='Crop'),
    gr.Number(label='Yield'),
]
outputs_traditional = gr.Textbox(label='Predicted Profit')

inputs_deep_learning = [
    gr.Dropdown(choices=list(deep_learning_models.keys()) + ['Random Forest'], label='Model'),
    gr.File(label='Upload TIFF File')
]
outputs_deep_learning = gr.Image(label='Prediction Overlay')

with gr.Blocks() as demo:
    with gr.Tab("Traditional ML Models"):
        gr.Interface(
            fn=predict_traditional,
            inputs=inputs_traditional,
            outputs=outputs_traditional,
            title="Profit Prediction using Traditional ML Models"
        )
        
    with gr.Tab("Deep Learning Models"):
        gr.Interface(
            fn=lambda model_name, file: predict_deep_learning(model_name, file) if model_name != 'Random Forest' else predict_random_forest(file),
            inputs=inputs_deep_learning,
            outputs=outputs_deep_learning,
            title="Crop Yield Prediction using Deep Learning Models and Random Forest"
        )

demo.launch()