from detectree2model.predictions.predict import run_detectree2
from polygons_processing.postpprocess_detectree2 import postprocess
from generate_tree_images.generate_tree_images import generate_tree_images
from classification.classification_predict import classify
import os
import json
import gradio as gr
import rasterio
import geopandas as gpd
import plotly.graph_objects as go
import numpy as np
import ast

def row_to_feature(row):
    feature = {
        "id": row["id"],
        "type": "Feature",
        "properties": {"Confidence_score": row["Confidence_score"], "species": row['species']},
        "geometry": {"type": "Polygon", "coordinates": [row["coordinates"]]},
    }
    return feature

def export_geojson(df, filename):
    features = [row_to_feature(row) for idx, row in df.iterrows()]

    feature_collection = {
        "type": "FeatureCollection",
        "crs": {"type": "name", "properties": {"name": "urn:ogc:def:crs:EPSG::32720"}},
        "features": features,
    }

    output_geojson = json.dumps(feature_collection)

    with open(f"{filename}.geojson", "w") as f:
        f.write(output_geojson)

    print(f"GeoJSON data exported to '{filename}.geojson' file.")

def process_image(image_path: str):
    current_directory = os.getcwd()

    output_directory = os.path.join(current_directory, "outputs")
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    run_detectree2(image_path, store_path=output_directory)

    processed_output_df = postprocess(output_directory + '/detectree2_delin.geojson', output_directory + '/processed_delin')

    processed_geojson = output_directory + '/processed_delin.geojson'

    generate_tree_images(processed_geojson, image_path)

    output_folder = './tree_images'

    all_top_3_list = []  # Initialize an empty list to accumulate all top_3 lists

    for file_name in os.listdir(output_folder):
        file_path = os.path.join(output_folder, file_name)
        print(file_path)
        probs = classify(file_path)
        top_3 = probs.head(3)
        top_3_list = [[cls, prob] for cls, prob in top_3.items()]

        # Accumulate the top_3_list for each file
        all_top_3_list.append(top_3_list)

    # Assign the accumulated top_3_list to the 'species' column of the dataframe
    processed_output_df['species'] = all_top_3_list

    final_output_path = 'result'

    export_geojson(processed_output_df, final_output_path)

    return final_output_path, image_path

def plot_results(geojson_path, image_path):
    with rasterio.open(image_path) as src:
        tif_image = src.read([1, 2, 3])  # Read the first three bands (RGB)
        tif_transform = src.transform
        height, width = tif_image.shape[1], tif_image.shape[2]

    # Read the GeoJSON file
    geojson_data = gpd.read_file(geojson_path + '.geojson')

    # Create Plotly figure
    fig = go.Figure()

    # Add image to the figure
    fig.add_trace(go.Image(z=tif_image.transpose((1, 2, 0)), hoverinfo = 'none'))

    # Add polygons to the plot
    for idx, row in geojson_data.iterrows():
        coordinates = row['geometry'].exterior.coords.xy
        x, y = list(coordinates[0]), list(coordinates[1])  # Convert to list

        # Transform coordinates to match image pixel space
        x_transformed = [(xi - tif_transform.c) / tif_transform.a for xi in x]
        y_transformed = [(yi - tif_transform.f) / tif_transform.e for yi in y]

        species_info_str = row['species']
        species_info = ast.literal_eval(species_info_str)
        first_array = species_info[0]
        second_array = species_info[1]
        third_array = species_info[2]
        confidence_score = row['Confidence_score']
        hovertemplate = f"Polygon:<br>{idx}<br><br>Species and Probability:<br>{first_array}<br>{second_array}<br>{third_array}<br><br>Confidence:<br>{confidence_score}"

        fig.add_trace(go.Scatter(
            x=x_transformed, 
            y=y_transformed, 
            mode='lines',  
            name = '',
            line=dict(color='red'),
            hovertemplate=hovertemplate,
            hoverinfo='text'
        ))

    fig.update_layout(
        title='TIF Image with Tree Crowns Overlay', 
        xaxis_title='X',
        yaxis_title='Y',
        showlegend=False, # Hide the legend
    )

    return fig



def greet(image_path: str):
    geojson_path, image_path = process_image(image_path)
    fig = plot_results(geojson_path, image_path)
    
    return fig
    
#tif_file_name = "TreeCrownVectorDataset_761588_9673769_20_20_32720.tif"
#tif_input = "/Users/jonathanseele/ETH/Hackathons/EcoHackathon/WeCanopy/test/" + tif_file_name

# File paths
#tif_file_path = '/Users/taekim/ecohackathon/WeCanopy/test/TreeCrownVectorDataset_761588_9673769_20_20_32720.tif'
#geojson_file_path = '/Users/taekim/ecohackathon/WeCanopy/test/result.geojson'

# Read the TIF file


def main():
    demo = gr.Interface(
        fn=greet,
        inputs=gr.File(type='filepath'),
        outputs=gr.Plot(label="Tree Crowns"),
    )

    demo.launch(server_name='0.0.0.0', debug=True)

if __name__ == "__main__":
    main()