import csv
import json
import os
import math
from collections import defaultdict, Counter
from datetime import datetime

class DataEngine:
    def __init__(self, csv_path: str):
        self.csv_path = csv_path
        self.records = []
        
        # Bounding box of Bengaluru violations (to be computed dynamically)
        self.min_lat = 90.0
        self.max_lat = -90.0
        self.min_lon = 180.0
        self.max_lon = -180.0
        
        # Grid settings
        self.grid_size = 100 # 100x100 grid cells
        
        # Overall statistics cached
        self.stats = {}
        
        # Load and process dataset
        self._load_data()

    def _load_data(self):
        print(f"Loading dataset from {self.csv_path}...")
        start_time = datetime.now()
        
        if not os.path.exists(self.csv_path):
            raise FileNotFoundError(f"Dataset file not found at: {self.csv_path}")

        temp_records = []
        with open(self.csv_path, 'r', encoding='utf-8', errors='ignore') as f:
            reader = csv.reader(f)
            headers = next(reader)
            
            # Find column indices
            lat_idx = headers.index('latitude')
            lon_idx = headers.index('longitude')
            v_type_idx = headers.index('vehicle_type')
            up_v_type_idx = headers.index('updated_vehicle_type')
            violation_idx = headers.index('violation_type')
            time_idx = headers.index('created_datetime')
            ps_idx = headers.index('police_station')
            
            for row in reader:
                try:
                    lat = float(row[lat_idx])
                    lon = float(row[lon_idx])
                    
                    # Skip rows with invalid or zero coordinates
                    if lat == 0.0 or lon == 0.0:
                        continue
                        
                    # Track bounds
                    if lat < self.min_lat: self.min_lat = lat
                    if lat > self.max_lat: self.max_lat = lat
                    if lon < self.min_lon: self.min_lon = lon
                    if lon > self.max_lon: self.max_lon = lon
                    
                    # Vehicle Type fallback
                    v_type = row[up_v_type_idx]
                    if v_type == 'NULL' or not v_type:
                        v_type = row[v_type_idx]
                    
                    # Parse violation list
                    violation_str = row[violation_idx]
                    try:
                        violations = json.loads(violation_str)
                    except:
                        # Fallback for irregular json format
                        violations = [v.strip('[]"\' ') for v in violation_str.split(',')]
                        
                    # Parse timestamp (e.g. 2023-11-20 00:28:46+00)
                    time_str = row[time_idx].split('+')[0].strip()
                    dt = datetime.strptime(time_str, "%Y-%m-%d %H:%M:%S")
                    
                    police_station = row[ps_idx]
                    if police_station == 'NULL':
                        police_station = 'Unknown'
                        
                    temp_records.append({
                        'lat': lat,
                        'lon': lon,
                        'vehicle_type': v_type,
                        'violations': violations,
                        'hour': dt.hour,
                        'day_of_week': dt.weekday(), # 0=Monday, 6=Sunday
                        'month': dt.month,
                        'police_station': police_station
                    })
                except Exception as e:
                    # Skip corrupted rows silently
                    continue
        
        # Apply a sanity padding to bounds to prevent out-of-bounds calculations
        self.min_lat -= 0.005
        self.max_lat += 0.005
        self.min_lon -= 0.005
        self.max_lon += 0.005
        
        self.records = temp_records
        duration = (datetime.now() - start_time).total_seconds()
        print(f"Loaded {len(self.records)} valid records in {duration:.2f} seconds.")
        print(f"Coordinates Bounds: Lat({self.min_lat:.4f} to {self.max_lat:.4f}), Lon({self.min_lon:.4f} to {self.max_lon:.4f})")
        
        # Precompute overall statistics
        self._precompute_stats()

    def get_grid_indices(self, lat: float, lon: float) -> tuple:
        """Map latitude and longitude to grid cell indices (x, y)."""
        x = int((lon - self.min_lon) / (self.max_lon - self.min_lon) * self.grid_size)
        y = int((lat - self.min_lat) / (self.max_lat - self.min_lat) * self.grid_size)
        
        # Clamp to bounds
        x = max(0, min(self.grid_size - 1, x))
        y = max(0, min(self.grid_size - 1, y))
        return (x, y)

    def get_grid_coordinates(self, x: int, y: int) -> tuple:
        """Get latitude and longitude of the center of a grid cell (x, y)."""
        cell_lon_width = (self.max_lon - self.min_lon) / self.grid_size
        cell_lat_height = (self.max_lat - self.min_lat) / self.grid_size
        
        lon = self.min_lon + (x + 0.5) * cell_lon_width
        lat = self.min_lat + (y + 0.5) * cell_lat_height
        return (lat, lon)

    def _get_vehicle_weight(self, v_type: str) -> float:
        v_type = v_type.upper()
        if any(keyword in v_type for keyword in ['BUS', 'TRUCK', 'HEAVY', 'CONSTRUCTION', 'STAGE CARRIAGE']):
            return 3.0
        if any(keyword in v_type for keyword in ['MAXI-CAB', 'LGV', 'VAN', 'TRACTOR']):
            return 2.0
        if any(keyword in v_type for keyword in ['CAR', 'GOODS AUTO', 'PASSENGER AUTO', 'AUTO']):
            return 1.5
        if any(keyword in v_type for keyword in ['SCOOTER', 'MOTOR CYCLE', 'MOPED', 'TWO WHEELER']):
            return 0.5
        return 1.0

    def _get_violation_weight(self, violations: list) -> float:
        max_weight = 0.5
        for v in violations:
            v = v.upper()
            if any(keyword in v for keyword in ['DOUBLE PARKING', 'TRAFFIC LIGHT', 'ZEBRA CROSS']):
                weight = 2.5
            elif any(keyword in v for keyword in ['WRONG PARKING', 'MAIN ROAD', 'FOOTPATH']):
                weight = 2.0
            elif any(keyword in v for keyword in ['NO PARKING', 'BUSTOP', 'SCHOOL', 'ROAD CROSSING']):
                weight = 1.5
            else:
                weight = 0.5
            max_weight = max(max_weight, weight)
        return max_weight

    def _get_priority_tier(self, gci: float) -> str:
        if gci >= 40.0:
            return "Critical"
        if gci >= 15.0:
            return "High"
        if gci >= 5.0:
            return "Watch"
        return "Low"

    def _build_explanation(self, records: list, gci: float, count: int) -> dict:
        vehicle_counter = Counter()
        violation_counter = Counter()
        hour_counter = Counter()
        peak_weighted = 0
        evidence_count = len(records)

        for r in records:
            vehicle_counter[r['vehicle_type']] += 1
            hour_counter[r['hour']] += 1
            if 8 <= r['hour'] <= 11 or 17 <= r['hour'] <= 20:
                peak_weighted += 1
            for v in r['violations']:
                violation_counter[v] += 1

        top_vehicle = vehicle_counter.most_common(1)[0][0] if vehicle_counter else "Unknown"
        top_violation = violation_counter.most_common(1)[0][0] if violation_counter else "Unknown"
        peak_hour = hour_counter.most_common(1)[0][0] if hour_counter else None
        peak_share = round((peak_weighted / max(1, count)) * 100.0, 1)
        tier = self._get_priority_tier(gci)
        if evidence_count >= 25:
            confidence = "High"
        elif evidence_count >= 8:
            confidence = "Medium"
        else:
            confidence = "Low"

        reasons = []
        if top_violation != "Unknown":
            reasons.append(f"{top_violation} is the dominant violation")
        if top_vehicle != "Unknown":
            reasons.append(f"{top_vehicle} appears most often")
        if peak_share >= 35.0:
            reasons.append(f"{peak_share}% of records fall in peak traffic windows")
        if count >= 10:
            reasons.append("high repeat frequency in this cell")

        return {
            'priority_tier': tier,
            'top_vehicle_type': top_vehicle,
            'top_violation_type': top_violation,
            'peak_hour': peak_hour,
            'peak_window_share_pct': peak_share,
            'evidence_records': evidence_count,
            'confidence': confidence,
            'reason': "; ".join(reasons[:3]) if reasons else "limited historical evidence, monitor for recurrence"
        }

    def attach_grid_explanations(self, grid_data: dict, day_of_week: int = None, hour: int = None, police_station: str = None) -> dict:
        """Attach explainability metadata to existing grid cells using historical records."""
        records_by_grid = defaultdict(list)
        fallback_records_by_grid = defaultdict(list)

        for r in self.records:
            if police_station is not None and police_station != 'All Stations' and r['police_station'] != police_station:
                continue

            x, y = self.get_grid_indices(r['lat'], r['lon'])
            grid_key = f"{x}_{y}"
            fallback_records_by_grid[grid_key].append(r)

            if day_of_week is not None and r['day_of_week'] != day_of_week:
                continue
            if hour is not None and r['hour'] != hour:
                continue

            records_by_grid[grid_key].append(r)

        for key, cell in grid_data.items():
            records = records_by_grid.get(key, []) or fallback_records_by_grid.get(key, [])
            cell['explanation'] = self._build_explanation(records, cell.get('gci', 0.0), cell.get('count', len(records)))

        return grid_data

    def _precompute_stats(self):
        """Precompute overall metadata to feed static API endpoints."""
        total = len(self.records)
        vehicle_counter = Counter()
        violation_counter = Counter()
        ps_counter = Counter()
        
        for r in self.records:
            vehicle_counter[r['vehicle_type']] += 1
            ps_counter[r['police_station']] += 1
            for v in r['violations']:
                violation_counter[v] += 1
                
        # Format stats
        self.stats = {
            'total_violations': total,
            'top_vehicle_types': [{'type': k, 'count': v, 'pct': round(v/total*100, 1)} for k, v in vehicle_counter.most_common(10)],
            'top_violation_types': [{'violation': k, 'count': v, 'pct': round(v/total*100, 1)} for k, v in violation_counter.most_common(10)],
            'top_police_stations': [{'station': k, 'count': v, 'pct': round(v/total*100, 1)} for k, v in ps_counter.most_common(10)],
            'bounds': {
                'min_lat': self.min_lat,
                'max_lat': self.max_lat,
                'min_lon': self.min_lon,
                'max_lon': self.max_lon
            }
        }

    def compute_grid_gci(self, day_of_week: int = None, hour: int = None, police_station: str = None) -> dict:
        """
        Calculate GCI for all grid cells, optionally filtered by day, hour, and police station.
        Returns a dictionary of grid_key: gci_score.
        """
        grid_scores = defaultdict(float)
        grid_counts = defaultdict(int)
        
        for r in self.records:
            # Apply filters
            if day_of_week is not None and r['day_of_week'] != day_of_week:
                continue
            if hour is not None and r['hour'] != hour:
                continue
            if police_station is not None and police_station != 'All Stations' and r['police_station'] != police_station:
                continue
                
            # Calculate weight
            v_weight = self._get_vehicle_weight(r['vehicle_type'])
            severity = self._get_violation_weight(r['violations'])
            
            # Temporal weight: Peak traffic hours (8-11 AM, 5-8 PM)
            time_factor = 1.5 if (8 <= r['hour'] <= 11 or 17 <= r['hour'] <= 20) else 1.0
            
            violation_gci = v_weight * severity * time_factor
            
            # Grid mapping
            x, y = self.get_grid_indices(r['lat'], r['lon'])
            grid_key = f"{x}_{y}"
            
            grid_scores[grid_key] += violation_gci
            grid_counts[grid_key] += 1
            
        # Structure grid output
        grid_data = {}
        for key, score in grid_scores.items():
            x, y = map(int, key.split('_'))
            lat, lon = self.get_grid_coordinates(x, y)
            grid_data[key] = {
                'x': x,
                'y': y,
                'lat': round(lat, 6),
                'lon': round(lon, 6),
                'gci': round(score, 1),
                'count': grid_counts[key]
            }
            
        return self.attach_grid_explanations(grid_data, day_of_week, hour, police_station)

    def compute_road_hotspots(self, day_of_week: int = None, hour: int = None, police_station: str = None, limit: int = 900) -> dict:
        """
        Aggregate violations at near-real coordinates for road-shaped map rendering.
        This avoids drawing artificial square grids over the road map.
        """
        buckets = {}

        for r in self.records:
            if day_of_week is not None and r['day_of_week'] != day_of_week:
                continue
            if hour is not None and r['hour'] != hour:
                continue
            if police_station is not None and police_station != 'All Stations' and r['police_station'] != police_station:
                continue

            # Roughly 10-12m buckets, enough to merge repeated violations on the same road edge.
            key = f"{round(r['lat'], 4)}_{round(r['lon'], 4)}"
            if key not in buckets:
                buckets[key] = {
                    'lat_sum': 0.0,
                    'lon_sum': 0.0,
                    'gci': 0.0,
                    'count': 0,
                    'records': []
                }

            v_weight = self._get_vehicle_weight(r['vehicle_type'])
            severity = self._get_violation_weight(r['violations'])
            time_factor = 1.5 if (8 <= r['hour'] <= 11 or 17 <= r['hour'] <= 20) else 1.0

            buckets[key]['lat_sum'] += r['lat']
            buckets[key]['lon_sum'] += r['lon']
            buckets[key]['gci'] += v_weight * severity * time_factor
            buckets[key]['count'] += 1
            buckets[key]['records'].append(r)

        ranked = sorted(buckets.items(), key=lambda item: item[1]['gci'], reverse=True)[:limit]
        hotspots = {}

        for idx, (_, bucket) in enumerate(ranked):
            lat = bucket['lat_sum'] / bucket['count']
            lon = bucket['lon_sum'] / bucket['count']
            x, y = self.get_grid_indices(lat, lon)
            key = f"road_{idx}"
            gci = round(bucket['gci'], 1)
            hotspots[key] = {
                'id': key,
                'render_type': 'road_hotspot',
                'x': x,
                'y': y,
                'lat': round(lat, 6),
                'lon': round(lon, 6),
                'gci': gci,
                'count': bucket['count'],
                'explanation': self._build_explanation(bucket['records'], gci, bucket['count'])
            }

        return hotspots

    def rank_police_stations(self, day_of_week: int = None, hour: int = None, limit: int = 10) -> list:
        """Rank stations by current hotspot burden using only historical violation data."""
        station_scores = defaultdict(float)
        station_counts = defaultdict(int)
        station_hours = defaultdict(Counter)
        station_violations = defaultdict(Counter)

        for r in self.records:
            if r['police_station'] == 'Unknown':
                continue
            if day_of_week is not None and r['day_of_week'] != day_of_week:
                continue
            if hour is not None and r['hour'] != hour:
                continue

            v_weight = self._get_vehicle_weight(r['vehicle_type'])
            severity = self._get_violation_weight(r['violations'])
            time_factor = 1.5 if (8 <= r['hour'] <= 11 or 17 <= r['hour'] <= 20) else 1.0
            score = v_weight * severity * time_factor
            station = r['police_station']

            station_scores[station] += score
            station_counts[station] += 1
            station_hours[station][r['hour']] += 1
            for v in r['violations']:
                station_violations[station][v] += 1

        rankings = []
        for station, score in station_scores.items():
            recommended_units = 1
            if score >= 600:
                recommended_units = 4
            elif score >= 300:
                recommended_units = 3
            elif score >= 120:
                recommended_units = 2

            top_violation = station_violations[station].most_common(1)
            peak_hour = station_hours[station].most_common(1)
            rankings.append({
                'station': station,
                'total_gci': round(score, 1),
                'records': station_counts[station],
                'peak_hour': peak_hour[0][0] if peak_hour else None,
                'top_violation': top_violation[0][0] if top_violation else 'Unknown',
                'recommended_units': recommended_units
            })

        rankings.sort(key=lambda row: row['total_gci'], reverse=True)
        return rankings[:limit]