from flask import Flask, render_template, request, redirect, url_for
import pandas as pd
import numpy as np
import joblib
import requests
from keras.models import model_from_json
import folium
import matplotlib.pyplot as plt
from io import BytesIO
import base64
import os
from collections import defaultdict

app = Flask(__name__)

# Load model and scalers
def load_model(name):
    with open(f"{name}.json", "r") as json_file:
        loaded_model_json = json_file.read()
    model = model_from_json(loaded_model_json)
    model.load_weights(f"{name}.weights.h5")
    return model

model = load_model("FUTURE_AQI_v1")
scaler_X = joblib.load('scaler_X_cpcb_4.pkl')
scaler_y = joblib.load('scaler_y_cpcb_4.pkl')

API_KEY = "26daca1b78f44099a755b921be4bfcf1"

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/forecast', methods=['POST'])
def forecast():
    # Get user input
    latitude = float(request.form['latitude'])
    longitude = float(request.form['longitude'])

    # Fetch current AQI from API
    current_url = f"https://api.weatherbit.io/v2.0/current/airquality?lat={latitude}&lon={longitude}&key={API_KEY}"
    response = requests.get(current_url)
    if response.status_code == 200:
        current_data = response.json()['data'][0]

    # Prepare input for the model
    now = pd.to_datetime("now")
    input_data = pd.DataFrame([{
        'PM2.5': current_data['pm25'],
        'PM10': current_data['pm10'],
        'NO2': current_data['no2'],
        'SO2': current_data['so2'],
        'CO': current_data['co'],
        'AQI': current_data['aqi'],
        'Day': now.day,
        'Month': now.month,
        'Hour': now.hour
    }])

    # Scale and predict
    input_scaled = scaler_X.transform(input_data)
    predictions = model.predict(input_scaled)
    predictions_actual = scaler_y.inverse_transform(predictions)

    # Fetch forecasted AQI from API
    forecast_url = f"https://api.weatherbit.io/v2.0/forecast/airquality?lat={latitude}&lon={longitude}&key={API_KEY}"
    response = requests.get(forecast_url)
    forecast_data = response.json()['data']

    grouped_aqi = defaultdict(list)
    for entry in forecast_data:
        date = entry['datetime'].split(':')[0]
        grouped_aqi[date].append(entry['aqi'])

    api_predictions = {date: max(values) for date, values in grouped_aqi.items()}

    # Save results to CSV
    forecast_df = pd.DataFrame([{
        **input_data.iloc[0],
        'lat': latitude,
        'lon': longitude,
        'AQI_step_1': predictions_actual[0, 0],
        'AQI_step_2': predictions_actual[0, 1],
        'AQI_step_3': predictions_actual[0, 2]
    }])
    forecast_df.to_csv('aqi_data.csv', mode='a', header=False, index=False)

    api_df = pd.DataFrame([{
        'AQI_currrent_API': current_data['aqi'],
        'AQI_step_1_API': api_predictions.get(list(api_predictions.keys())[0], None),
        'AQI_step_2_API': api_predictions.get(list(api_predictions.keys())[1], None),
        'AQI_step_3_API': api_predictions.get(list(api_predictions.keys())[2], None)
    }])
    api_df.to_csv('aqi_data_actual_api.csv', mode='a', header=False, index=False)

    # Generate updated map and return
    generate_map()
    return redirect(url_for('result'))

@app.route('/result')
def result():
    return render_template('result.html')

def generate_map():
    # Load data
    df1 = pd.read_csv('aqi_data.csv')
    df2 = pd.read_csv('aqi_data_actual_api.csv')
    data = pd.concat([df1, df2], axis=1)

    # Create Folium map
    # Create the Folium map
    map_center = [data['lat'].mean(), data['lon'].mean()]
    m = folium.Map(location=map_center, zoom_start=10)

    # AQI Color Legend
    legend_html = """
    <div style="
        position: fixed; 
        bottom: 20px; left: 20px; width: 350px; height: 225px; 
        background-color: white; 
        z-index:9999; font-size:14px; border:2px solid grey; 
        padding: 10px; overflow-y: auto;">
        <b>AQI Color Legend</b>
        <table style="width: 100%; border-collapse: collapse; text-align: left;">
            <thead>
                <tr style="border-bottom: 2px solid grey;">
                    <th style="padding: 5px;">Color</th>
                    <th style="padding: 5px;">Remark</th>
                    <th style="padding: 5px;">Range</th>
                </tr>
            </thead>
            <tbody>
                <tr>
                    <td><i style="background:green; width:15px; height:15px; display:inline-block; border:1px solid black;"></i></td>
                    <td>Good</td>
                    <td>0-50</td>
                </tr>
                <tr>
                    <td><i style="background:yellow; width:15px; height:15px; display:inline-block; border:1px solid black;"></i></td>
                    <td>Moderate</td>
                    <td>51-100</td>
                </tr>
                <tr>
                    <td><i style="background:orange; width:15px; height:15px; display:inline-block; border:1px solid black;"></i></td>
                    <td>Unhealthy for Sensitive Groups</td>
                    <td>101-150</td>
                </tr>
                <tr>
                    <td><i style="background:red; width:15px; height:15px; display:inline-block; border:1px solid black;"></i></td>
                    <td>Unhealthy</td>
                    <td>151-200</td>
                </tr>
                <tr>
                    <td><i style="background:purple; width:15px; height:15px; display:inline-block; border:1px solid black;"></i></td>
                    <td>Very Unhealthy</td>
                    <td>201-300</td>
                </tr>
                <tr>
                    <td><i style="background:maroon; width:15px; height:15px; display:inline-block; border:1px solid black;"></i></td>
                    <td>Hazardous</td>
                    <td>301+</td>
                </tr>
            </tbody>
        </table>
    </div>
    """
    m.get_root().html.add_child(folium.Element(legend_html))
    
    for _, row in data.iterrows():
        popup_html = create_plot(row)
        color = get_color_for_aqi(row['AQI_step_1'])
        folium.Marker(
            location=[row["lat"], row["lon"]],
            popup=folium.Popup(html=popup_html, max_width=500),
            icon=folium.Icon(color=color)
        ).add_to(m)

    # Save the map
    m.save('static/aqi_forecast_with_legend.html')

def create_plot(data):
    # Bar plot generation logic (same as before)
    fig, ax = plt.subplots(figsize=(5, 2))
    categories = ['DAY 1', 'DAY 2', 'DAY 3']
    actual_values = [data['AQI_step_1'], data['AQI_step_2'], data['AQI_step_3']]
    api_values = [data['AQI_step_1_API'], data['AQI_step_2_API'], data['AQI_step_3_API']]

    bar_width = 0.35
    index = range(len(categories))

    # Plot horizontal bars
    bars_actual = ax.barh(index, actual_values, bar_width, label="Model AQI", color='blue')
    bars_api = ax.barh([i + bar_width for i in index], api_values, bar_width, label="API AQI", color='green')

    # Add values to each bar
    max_value = 0  # Track the maximum value for axis limit adjustment
    for bar in bars_actual:
        value = bar.get_width()
        ax.text(value + 2, bar.get_y() + bar.get_height() / 2, 
                f'{value:.1f}', va='center', fontsize=10)
        max_value = max(max_value, value)
    for bar in bars_api:
        value = bar.get_width()
        ax.text(value + 2, bar.get_y() + bar.get_height() / 2, 
                f'{value:.1f}', va='center', fontsize=10)
        max_value = max(max_value, value)

    # Adjust x-axis limits to accommodate annotations
    ax.set_xlim(0, max_value * 1.2)

    # Customize y-ticks and labels
    ax.set_yticks([i + bar_width / 2 for i in index])
    ax.set_yticklabels(categories)
    ax.set_xlabel('AQI')
    ax.set_title('AQI Comparison')

    # Place legend outside the plot area
    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), frameon=False)

    plt.tight_layout()

    # Save the plot to a PNG image in memory
    buffer = BytesIO()
    plt.savefig(buffer, format="png", bbox_inches='tight')
    plt.close(fig)
    buffer.seek(0)

    # Encode the image to base64 to embed it in the HTML
    image_base64 = base64.b64encode(buffer.read()).decode()
    return f'<img src="data:image/png;base64,{image_base64}">'

def get_color_for_aqi(aqi_value):
    # Color logic (same as before)
    if aqi_value <= 50:
        return 'green'
    elif aqi_value <= 100:
        return 'yellow'
    elif aqi_value <= 150:
        return 'orange'
    elif aqi_value <= 200:
        return 'red'
    elif aqi_value <= 300:
        return 'purple'
    else:
        return 'maroon'

if __name__ == '__main__':
    app.run(debug=True)