from flask import Flask, render_template, jsonify
import folium
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from io import BytesIO
import base64

# Zone Pressure Monitor
class ZonePressureMonitor:
    def __init__(self):
        self.pressures = {}

    def update_pressure(self, zone, pressure):
        self.pressures[zone] = pressure

    def get_pressures(self):
        return self.pressures

# Heatmap Generator
class HeatmapGenerator:
    def __init__(self, pressure_data):
        self.pressure_data = pressure_data

    def generate_heatmap(self):
        zones = list(self.pressure_data.keys())
        pressures = list(self.pressure_data.values())
        zone_grid = np.array(pressures).reshape(len(zones), 1)

        plt.figure(figsize=(6, 4))
        sns.heatmap(zone_grid, annot=True, xticklabels=zones, yticklabels=["Pressure"], cmap='YlGnBu')
        plt.title('Zone Pressure Heatmap')
        plt.xlabel('Zones')
        plt.ylabel('Pressure Levels')
        
        buf = BytesIO()
        plt.savefig(buf, format="png")
        plt.close()
        return base64.b64encode(buf.getvalue()).decode('utf-8')

# Reward System
class RewardSystem:
    def __init__(self):
        self.units = {}

    def record_efficiency(self, unit_id, efficiency):
        self.units[unit_id] = efficiency

    def calculate_rewards(self):
        rewards = {}
        for unit, efficiency in self.units.items():
            if efficiency >= 75:
                rewards[unit] = 100
            elif efficiency >= 50:
                rewards[unit] = 50
            else:
                rewards[unit] = 0
        return rewards

# Driver Assignment
class DriverAssignment:
    def __init__(self, zone_pressures):
        self.zone_pressures = zone_pressures

    def optimize_assignments(self, drivers):
        assignments = {}
        sorted_zones = sorted(self.zone_pressures.items(), key=lambda item: item[1])
        for i, driver in enumerate(drivers):
            if i < len(sorted_zones):
                assignments[driver] = sorted_zones[i][0]
            else:
                assignments[driver] = 'No Zone Available'
        return assignments

app = Flask(__name__)

# Initialize components
zone_pressure_monitor = ZonePressureMonitor()
reward_system = RewardSystem()

@app.route('/pressure')
def pressure_dashboard():
    # Update and get zone pressures
    zone_pressure_monitor.update_pressure('Zone A', 30)
    zone_pressure_monitor.update_pressure('Zone B', 45)
    pressures = zone_pressure_monitor.get_pressures()

    # Generate heatmap
    heatmap_generator = HeatmapGenerator(pressures)
    heatmap_img = heatmap_generator.generate_heatmap()

    # Calculate rewards
    reward_system.record_efficiency('Vehicle1', 80)
    reward_system.record_efficiency('Vehicle2', 45)
    rewards = reward_system.calculate_rewards()

    # Optimize driver assignments
    drivers = ['Driver1', 'Driver2', 'Driver3']
    driver_assignment_system = DriverAssignment(pressures)
    assignments = driver_assignment_system.optimize_assignments(drivers)

    return render_template('pressure.html', heatmap_img=heatmap_img, rewards=rewards, assignments=assignments)

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