# td_traffic_spot_visualiser.py

# This module handles traffic data integration for the BNB+ platform, providing traffic-based
# discount calculations and map visualization of traffic spots. It includes classes for
# individual traffic spots and a manager to handle collections of spots.

# The module integrates with a dataset of traffic observations to determine traffic conditions
# and calculate eco-friendly discounts for BNB listings based on local traffic volume.

# Author: Gordon Li (20317033)
# Date: March 2025

import folium
import oracledb
import logging
import base64
import numpy as np
from html import escape
from datasets import load_dataset
from constant.hkust_bnb_constant import (
    GET_TRAFFIC_CAMERA_LOCATIONS,
    TRAFFIC_DISCOUNT_DISPLAY,
    TRAFFIC_POPUP_BASE,
    TRAFFIC_RECORDS_HEADER,
    TRAFFIC_RECORD_ENTRY,
    TRAFFIC_IMAGE_HTML,
    TRAFFIC_NO_RECORDS
)


class TDTrafficSpot:
    # Initializes a traffic spot with location and historical traffic data.

    # Parameters:
    #     key: Unique identifier for the traffic spot
    #     latitude: Geographic latitude of the traffic spot
    #     longitude: Geographic longitude of the traffic spot
    #     dataset_rows: List of dictionaries containing historical traffic observations (default: None)

    def __init__(self, key, latitude, longitude, dataset_rows=None):
        self.key = key
        self.latitude = float(latitude) if latitude is not None else None
        self.longitude = float(longitude) if longitude is not None else None
        self.dataset_rows = dataset_rows or []
        self.avg_vehicle_count = self.calculate_avg_vehicle_count()
        self.recent_display_rows = self.get_recent_display_rows()

    # Checks if the traffic spot has valid geographic coordinates.

    # Returns:
    #     Boolean indicating whether latitude and longitude are valid

    def is_valid(self):
        return self.latitude is not None and self.longitude is not None

    # Gets the most recent traffic observations for display purposes.

    # Parameters:
    #     max_display: Maximum number of recent records to return (default: 2)

    # Returns:
    #     List of the most recent traffic observation records

    def get_recent_display_rows(self, max_display=2):
        if not self.dataset_rows:
            return []

        sorted_rows = sorted(self.dataset_rows, key=lambda x: x['capture_time'], reverse=True)
        return sorted_rows[:max_display]

    # Calculates the average vehicle count based on historical traffic observations.

    # Returns:
    #     Float representing the average number of vehicles observed

    def calculate_avg_vehicle_count(self):
        if not self.dataset_rows:
            return 0

        vehicle_counts = [row.get('vehicle_count', 0) for row in self.dataset_rows if 'vehicle_count' in row]

        if not vehicle_counts:
            return 0

        return np.mean(vehicle_counts)

    # Determines the discount rate based on average traffic volume.

    # Returns:
    #     Float representing the discount rate (0.0 to 0.20)

    def get_discount_rate(self):
        if self.avg_vehicle_count < 2:
            return 0.20
        elif self.avg_vehicle_count <= 5:
            return 0.10
        else:
            return 0.0

    # Generates a human-readable description of the traffic-based discount.

    # Returns:
    #     String describing the discount, if any

    def get_discount_info(self):
        discount_rate = self.get_discount_rate()

        if discount_rate <= 0:
            return "No traffic discount available"

        return f"{int(discount_rate * 100)}% discount! Low traffic area"

    # Creates HTML content for the traffic spot's popup on the map.

    # Returns:
    #     HTML string for the Folium popup

    def create_popup_content(self):
        discount_info = self.get_discount_info()
        discount_display = TRAFFIC_DISCOUNT_DISPLAY.format(
                discount_info=discount_info,
                avg_vehicle_count=self.avg_vehicle_count,
                observation_count=len(self.dataset_rows)
            )

        html = TRAFFIC_POPUP_BASE.format(
            location_id=escape(str(self.key)),
            discount_display=discount_display
        )

        recent_rows = self.recent_display_rows

        if recent_rows:
            html += TRAFFIC_RECORDS_HEADER.format(
                recent_count=len(recent_rows),
                total_count=len(self.dataset_rows)
            )

            for row in recent_rows:
                image_data = row.get('processed_image')
                image_html = ""
                if image_data:
                    try:
                        base64_encoded = base64.b64encode(image_data).decode('utf-8')
                        image_html = TRAFFIC_IMAGE_HTML.format(base64_encoded=base64_encoded)
                    except Exception as e:
                        logging.error(f"Error encoding image for {self.key}: {str(e)}")
                        image_html = "<p>Image load failed</p>"

                html += TRAFFIC_RECORD_ENTRY.format(
                    capture_time=escape(str(row['capture_time'])),
                    vehicle_count=escape(str(row['vehicle_count'])),
                    image_html=image_html
                )
        else:
            html += TRAFFIC_NO_RECORDS

        html += "</div>"
        return html

    # Adds the traffic spot to a Folium map with appropriate styling.

    # Parameters:
    #     folium_map: Folium map object to add the marker to

    def add_to_map(self, folium_map):
        if self.is_valid():

            if self.avg_vehicle_count < 2:
                color = 'blue'  # Low traffic - 20% discount
            elif self.avg_vehicle_count < 5:
                color = 'orange'  # Medium traffic - 10% discount
            else:
                color = 'purple'  # High traffic - no discount

            folium.Marker(
                location=[self.latitude, self.longitude],
                popup=self.create_popup_content(),
                icon=folium.Icon(color=color, icon='camera'),
            ).add_to(folium_map)


class TrafficSpotManager:
    # Manages a collection of traffic spots, handling data loading and map integration.

    # Initializes the manager with database connection parameters and loads initial traffic spots.

    # Parameters:
    #     connection_params: Dictionary containing Oracle database connection parameters

    def __init__(self, connection_params):
        self.connection_params = connection_params
        self.traffic_spots = []
        self.spot_dict = {}
        self.load_limited_traffic_spots()

    # Loads a limited number of traffic spots for initial display.

    # Parameters:
    #     limit: Maximum number of traffic spots to load initially (default: 10)

    def load_limited_traffic_spots(self, limit=10):
        try:
            dataset = load_dataset("slliac/isom5240-td-application-traffic-analysis", split="application")
            dataset_list = list(dataset)

            location_data = {}
            for row in dataset_list:
                loc_id = row['location_id']
                if loc_id not in location_data:
                    location_data[loc_id] = []
                location_data[loc_id].append(row)

            if len(location_data) > limit:
                recent_activities = {}
                for loc_id, rows in location_data.items():
                    if rows:
                        most_recent = max(rows, key=lambda x: x['capture_time'])
                        recent_activities[loc_id] = most_recent['capture_time']

                top_locations = sorted(recent_activities.items(), key=lambda x: x[1], reverse=True)[:limit]
                selected_locations = [loc_id for loc_id, _ in top_locations]

                location_data = {loc_id: location_data[loc_id] for loc_id in selected_locations}

            if not location_data:
                logging.warning("No locations found in dataset")
                return

            location_ids = tuple(location_data.keys())

            with oracledb.connect(**self.connection_params) as conn:
                cursor = conn.cursor()

                placeholders = ','.join([':' + str(i + 1) for i in range(len(location_ids))])

                query = GET_TRAFFIC_CAMERA_LOCATIONS.format(placeholders=placeholders)

                cursor.execute(query, location_ids)
                spots = cursor.fetchall()

                self.traffic_spots = [
                    TDTrafficSpot(
                        spot[0],
                        spot[1],
                        spot[2],
                        location_data.get(spot[0], [])
                    )
                    for spot in spots
                ]

                for spot in self.traffic_spots:
                    self.spot_dict[spot.key] = spot

                logging.info(f"Loaded {len(self.traffic_spots)} traffic spots with full historical data")

        except Exception as e:
            logging.error(f"Error loading traffic spots: {str(e)}")
            self.traffic_spots = []
            self.spot_dict = {}

    # Loads specific traffic spots by their keys when needed.

    # Parameters:
    #     keys: List of traffic spot keys to load

    def load_specific_traffic_spots(self, keys):
        needed_keys = [key for key in keys if key not in self.spot_dict]

        if not needed_keys:
            return

        try:
            dataset = load_dataset("slliac/isom5240-td-application-traffic-analysis", split="application")
            dataset_list = list(dataset)

            location_data = {}
            for row in dataset_list:
                loc_id = row['location_id']
                if loc_id in needed_keys:
                    if loc_id not in location_data:
                        location_data[loc_id] = []
                    location_data[loc_id].append(row)

            if location_data and needed_keys:
                with oracledb.connect(**self.connection_params) as conn:
                    cursor = conn.cursor()

                    placeholders = ','.join([':' + str(i + 1) for i in range(len(needed_keys))])

                    query = GET_TRAFFIC_CAMERA_LOCATIONS.format(placeholders=placeholders)

                    cursor.execute(query, tuple(needed_keys))
                    spots = cursor.fetchall()

                    new_spots = [
                        TDTrafficSpot(
                            spot[0],
                            spot[1],
                            spot[2],
                            location_data.get(spot[0], [])
                        )
                        for spot in spots
                    ]

                    for spot in new_spots:
                        self.spot_dict[spot.key] = spot
                        self.traffic_spots.append(spot)

                    logging.info(f"Loaded {len(new_spots)} additional traffic spots with full historical data")

        except Exception as e:
            logging.error(f"Error loading specific traffic spots: {str(e)}")

    # Adds traffic spots to a Folium map.

    # Parameters:
    #     folium_map: Folium map object to add markers to
    #     spot_keys: Optional list of specific spot keys to add (default: None, adds all spots)

    def add_spots_to_map(self, folium_map, spot_keys=None):
        if spot_keys is None:
            for spot in self.traffic_spots:
                spot.add_to_map(folium_map)
        else:
            for key in spot_keys:
                if key in self.spot_dict:
                    self.spot_dict[key].add_to_map(folium_map)

    # Retrieves a traffic spot by its key, loading it if necessary.

    # Parameters:
    #     key: The unique identifier of the traffic spot

    # Returns:
    #     TDTrafficSpot object or None if not found

    def get_spot_by_key(self, key):
        if key in self.spot_dict:
            return self.spot_dict[key]

        self.load_specific_traffic_spots([key])
        return self.spot_dict.get(key)