# carbon_forecast.py
import requests
import pandas as pd
import numpy as np
from datetime import datetime, timedelta, timezone
from typing import Dict, List, Optional
from dataclasses import dataclass
import json
import time


@dataclass
class CarbonForecast:
    """Carbon intensity forecast for a specific region."""
    region: str
    timestamps: List[datetime]
    intensities: List[float]  # gCO2/kWh
    confidence_low: List[float]
    confidence_high: List[float]
    renewable_percentage: List[float]


class CarbonForecastProvider:
    """Fetches and forecasts carbon intensity data."""

    def __init__(self, electricitymap_token: Optional[str] = None):
        self.em_token = electricitymap_token
        self.base_url = "https://api.co2signal.com/v1"
        self._cache: Dict[str, tuple] = {}

    def get_current_intensity(self, zone_id: str) -> float:
        """Get current carbon intensity for a zone."""
        cache_key = f"current_{zone_id}"
        now = time.time()

        if cache_key in self._cache:
            cached_time, cached_val = self._cache[cache_key]
            if now - cached_time < 300:
                return cached_val

        headers = {}
        if self.em_token:
            headers["auth-token"] = self.em_token

        try:
            resp = requests.get(
                f"{self.base_url}/latest",
                headers=headers,
                params={"zone": zone_id},
                timeout=10,
            )
            resp.raise_for_status()
            data = resp.json()
            intensity = data["data"]["carbonIntensity"]
            self._cache[cache_key] = (now, intensity)
            return intensity
        except Exception as e:
            print(f"Warning: {e}, using default intensity")
            return self._get_default_intensity(zone_id)

    def get_forecast(
        self, zone_id: str, hours_ahead: int = 72
    ) -> CarbonForecast:
        """Get carbon intensity forecast for next N hours."""
        current = self.get_current_intensity(zone_id)

        now = datetime.now(timezone.utc)
        timestamps = [now + timedelta(hours=h) for h in range(hours_ahead)]

        intensities = self._generate_forecast(current, hours_ahead, zone_id)

        noise = np.random.normal(0, 0.1, hours_ahead)
        confidence_low = [max(0, i * (1 + n - 0.15)) for i, n in zip(intensities, noise)]
        confidence_high = [i * (1 + abs(n) + 0.15) for i, n in zip(intensities, noise)]

        renewable = [self._estimate_renewable_pct(i) for i in intensities]

        return CarbonForecast(
            region=zone_id,
            timestamps=timestamps,
            intensities=intensities,
            confidence_low=confidence_low,
            confidence_high=confidence_high,
            renewable_percentage=renewable,
        )

    def _generate_forecast(
        self, current: float, hours: int, zone_id: str
    ) -> List[float]:
        """Generate synthetic forecast based on historical patterns."""
        hour_of_day = datetime.now(timezone.utc).hour
        forecast = []

        for h in range(hours):
            future_hour = (hour_of_day + h) % 24

            solar_factor = max(0, np.sin((future_hour - 6) * np.pi / 12))
            wind_factor = 0.3 + 0.2 * np.sin((future_hour + 6) * np.pi / 12)
            demand_factor = 1 + 0.3 * np.sin((future_hour - 14) * np.pi / 12)

            renewable_pct = 0.4 * solar_factor + 0.3 * wind_factor + 0.3
            intensity = current * (1 + 0.3 * (demand_factor - 1)) * (1 - 0.5 * renewable_pct)

            intensity = max(20, min(900, intensity + np.random.normal(0, 15)))
            forecast.append(round(intensity, 1))

        return forecast

    def _estimate_renewable_pct(self, intensity: float) -> float:
        """Estimate renewable percentage from intensity."""
        return max(0, min(100, 100 - (intensity / 8)))

    def _get_default_intensity(self, zone_id: str) -> float:
        """Default intensity values."""
        defaults = {
            "DE": 338, "FR": 56, "GB": 220, "NO": 10,
            "SE": 13, "US-CAL-CISO": 230, "US-MIDA-PJM": 350,
            "US-NW-PACW": 120, "US-TEX-ERCO": 450,
        }
        return defaults.get(zone_id, 400)

# carbon_scheduler.py
import pulp
import numpy as np
from typing import List, Dict, Optional
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from carbon_forecast import CarbonForecast


@dataclass
class TrainingJob:
    """Represents an ML training job."""
    job_id: str
    gpu_hours_required: float
    gpu_type: str = "A100"
    deadline_hours: float = 168  # 7 days default
    priority: int = 5  # 1 (lowest) to 10 (highest)
    preemptible: bool = True
    checkpoint_interval_minutes: int = 30
    region_preference: Optional[str] = None
    max_carbon_kg: Optional[float] = None
    max_cost_usd: Optional[float] = None
    estimated_carbon_kg: float = 0.0
    estimated_cost_usd: float = 0.0


@dataclass
class ScheduleSlot:
    """A scheduled time slot for a job."""
    job_id: str
    region: str
    start_time: datetime
    end_time: datetime
    gpu_count: int
    intensity_at_start: float
    estimated_emissions_kg: float
    estimated_cost_usd: float


class CarbonAwareScheduler:
    """Optimizes training schedules for minimal carbon emissions."""

    def __init__(self, forecast_provider):
        self.forecast_provider = forecast_provider
        self.scheduled_jobs: Dict[str, List[ScheduleSlot]] = {}

    def schedule_jobs(
        self,
        jobs: List[TrainingJob],
        regions: List[str] = ["DE", "FR", "NO", "SE"],
        max_gpus_per_region: int = 64,
    ) -> Dict[str, List[ScheduleSlot]]:
        """Schedule jobs to minimize carbon within constraints."""
        forecasts = {}
        for region in regions:
            forecasts[region] = self.forecast_provider.get_forecast(region)

        problem = pulp.LpProblem("CarbonAwareScheduling", pulp.LpMinimize)

        job_region_time = {}
        for job in jobs:
            job_region_time[job.job_id] = {}
            for region in regions:
                forecast = forecasts[region]
                n_slots = min(
                    len(forecast.timestamps),
                    int(job.deadline_hours),
                )
                job_region_time[job.job_id][region] = list(range(n_slots))

        x = {}
        for job in jobs:
            for region in regions:
                for t in job_region_time[job.job_id][region]:
                    x[(job.job_id, region, t)] = pulp.LpVariable(
                        f"x_{job.job_id}_{region}_{t}", cat="Binary"
                    )

        objective = pulp.lpSum(
            x[(job.job_id, region, t)]
            * forecasts[region].intensities[t]
            * job.gpu_hours_required
            * 0.3  # kWh per GPU-hour
            / 1000  # Convert to kg
            for job in jobs
            for region in regions
            for t in job_region_time[job.job_id][region]
        )
        problem += objective

        for job in jobs:
            problem += (
                pulp.lpSum(
                    x[(job.job_id, region, t)]
                    for region in regions
                    for t in job_region_time[job.job_id][region]
                ) == 1,
                f"assign_{job.job_id}",
            )

        for region in regions:
            forecast = forecasts[region]
            for t in range(len(forecast.timestamps)):
                problem += (
                    pulp.lpSum(
                        x[(job.job_id, region, t)]
                        for job in jobs
                        if t in job_region_time.get(job.job_id, {}).get(region, [])
                    ) <= max_gpus_per_region,
                    f"capacity_{region}_{t}",
                )

        for job in jobs:
            if job.max_carbon_kg:
                problem += (
                    pulp.lpSum(
                        x[(job.job_id, region, t)]
                        * forecasts[region].intensities[t]
                        * job.gpu_hours_required
                        * 0.3 / 1000
                        for region in regions
                        for t in job_region_time[job.job_id][region]
                    ) <= job.max_carbon_kg,
                    f"carbon_{job.job_id}",
                )

        solver = pulp.PULP_CBC_CMD(msg=0)
        problem.solve(solver)

        if problem.status != pulp.constants.LpStatusOptimal:
            print("Warning: No optimal solution found, using greedy fallback")
            return self._greedy_schedule(jobs, regions, forecasts)

        schedule = {}
        for job in jobs:
            schedule[job.job_id] = []
            for region in regions:
                for t in job_region_time[job.job_id][region]:
                    if x[(job.job_id, region, t)].varValue > 0.5:
                        forecast = forecasts[region]
                        start_time = forecast.timestamps[t]
                        end_time = start_time + timedelta(
                            hours=job.gpu_hours_required
                        )
                        emissions = (
                            forecast.intensities[t]
                            * job.gpu_hours_required
                            * 0.3 / 1000
                        )
                        slot = ScheduleSlot(
                            job_id=job.job_id,
                            region=region,
                            start_time=start_time,
                            end_time=end_time,
                            gpu_count=1,
                            intensity_at_start=forecast.intensities[t],
                            estimated_emissions_kg=round(emissions, 4),
                            estimated_cost_usd=round(
                                job.gpu_hours_required * self._get_region_cost(region), 2
                            ),
                        )
                        schedule[job.job_id].append(slot)

        self.scheduled_jobs = schedule
        return schedule

    def _greedy_schedule(
        self,
        jobs: List[TrainingJob],
        regions: List[str],
        forecasts: Dict[str, CarbonForecast],
    ) -> Dict[str, List[ScheduleSlot]]:
        """Fallback greedy scheduler when ILP fails."""
        schedule = {}
        sorted_jobs = sorted(jobs, key=lambda j: -j.priority)

        for job in sorted_jobs:
            best_region = None
            best_intensity = float("inf")
            best_time = 0

            for region in regions:
                forecast = forecasts[region]
                n_slots = min(len(forecast.intensities), int(job.deadline_hours))
                for t in range(n_slots):
                    if forecast.intensities[t] < best_intensity:
                        best_intensity = forecast.intensities[t]
                        best_region = region
                        best_time = t

            if best_region:
                forecast = forecasts[best_region]
                start = forecast.timestamps[best_time]
                end = start + timedelta(hours=job.gpu_hours_required)
                emissions = best_intensity * job.gpu_hours_required * 0.3 / 1000
                schedule[job.job_id] = [ScheduleSlot(
                    job_id=job.job_id,
                    region=best_region,
                    start_time=start,
                    end_time=end,
                    gpu_count=1,
                    intensity_at_start=best_intensity,
                    estimated_emissions_kg=round(emissions, 4),
                    estimated_cost_usd=round(
                        job.gpu_hours_required * self._get_region_cost(best_region), 2
                    ),
                )]

        return schedule

    @staticmethod
    def _get_region_cost(region: str) -> float:
        """Get approximate GPU-hour cost by region (USD)."""
        costs = {
            "DE": 2.50, "FR": 2.30, "NO": 2.80, "SE": 2.70,
            "US-CAL-CISO": 3.00, "US-MIDA-PJM": 2.80,
            "US-NW-PACW": 2.60, "US-TEX-ERCO": 2.40,
        }
        return costs.get(region, 2.50)

    def compare_naive_vs_aware(
        self, job: TrainingJob, regions: List[str]
    ) -> Dict:
        """Compare naive scheduling vs carbon-aware scheduling."""
        naive_emissions = 0
        aware_emissions = 0

        for region in regions:
            forecast = self.forecast_provider.get_forecast(region)
            if forecast.intensities:
                naive_emissions += np.mean(forecast.intensities) * job.gpu_hours_required * 0.3 / 1000
                min_idx = np.argmin(forecast.intensities[:int(job.deadline_hours)])
                aware_emissions += forecast.intensities[min_idx] * job.gpu_hours_required * 0.3 / 1000

        naive_emissions /= len(regions)
        aware_emissions /= len(regions)

        return {
            "naive_emissions_kg": round(naive_emissions, 4),
            "aware_emissions_kg": round(aware_emissions, 4),
            "reduction_kg": round(naive_emissions - aware_emissions, 4),
            "reduction_percent": round(
                (1 - aware_emissions / naive_emissions) * 100, 1
            ) if naive_emissions > 0 else 0,
        }

# checkpoint_manager.py
import torch
import os
import json
from datetime import datetime
from typing import Optional
from pathlib import Path


class CarbonAwareCheckpointManager:
    """Manages checkpoints for preemption-resilient training."""

    def __init__(
        self,
        checkpoint_dir: str = "./checkpoints",
        max_checkpoints: int = 3,
    ):
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.max_checkpoints = max_checkpoints
        self.checkpoint_index = 0

    def save_checkpoint(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        scheduler,
        epoch: int,
        step: int,
        metrics: dict,
        carbon_data: dict,
    ) -> str:
        """Save training checkpoint with carbon tracking data."""
        checkpoint = {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
            "epoch": epoch,
            "step": step,
            "metrics": metrics,
            "carbon_data": carbon_data,
            "timestamp": datetime.now().isoformat(),
            "checkpoint_index": self.checkpoint_index,
        }

        path = self.checkpoint_dir / f"checkpoint_{self.checkpoint_index:06d}.pt"
        torch.save(checkpoint, path)

        self._cleanup_old_checkpoints()
        self.checkpoint_index += 1

        return str(path)

    def load_checkpoint(
        self, path: Optional[str] = None
    ) -> Optional[dict]:
        """Load the latest or specified checkpoint."""
        if path:
            return torch.load(path) if os.path.exists(path) else None

        checkpoints = sorted(self.checkpoint_dir.glob("checkpoint_*.pt"))
        if not checkpoints:
            return None

        latest = checkpoints[-1]
        return torch.load(latest)

    def _cleanup_old_checkpoints(self):
        """Remove old checkpoints, keeping only max_checkpoints."""
        checkpoints = sorted(self.checkpoint_dir.glob("checkpoint_*.pt"))
        while len(checkpoints) > self.max_checkpoints:
            oldest = checkpoints.pop(0)
            oldest.unlink()

    def get_resume_info(self) -> dict:
        """Get information needed to resume training."""
        checkpoint = self.load_checkpoint()
        if checkpoint is None:
            return {"can_resume": False}

        return {
            "can_resume": True,
            "epoch": checkpoint["epoch"],
            "step": checkpoint["step"],
            "metrics": checkpoint["metrics"],
            "carbon_data": checkpoint.get("carbon_data", {}),
            "timestamp": checkpoint.get("timestamp", ""),
        }