# carbon_calculator.py
import time
import json
import logging
from datetime import datetime, timezone
from typing import Optional, Dict, Any
from dataclasses import dataclass, field
import requests
from codecarbon import EmissionsTracker

logger = logging.getLogger(__name__)


@dataclass
class CarbonConfig:
    """Configuration for carbon footprint calculation."""
    project_name: str
    experiment_name: str
    team_name: str = "default"
    country_code: str = "USA"
    cloud_provider: Optional[str] = None
    cloud_region: Optional[str] = None
    pue_factor: float = 1.2
    gpu_count: int = 1
    gpu_type: str = "NVIDIA A100"
    enable_db_logging: bool = True
    carbon_offset_per_kg: float = 0.0
    monthly_budget_kg: float = 1000.0


@dataclass
class CarbonReport:
    """Complete carbon emission report."""
    experiment_name: str
    model_name: str
    duration_hours: float
    total_energy_kwh: float
    gpu_energy_kwh: float
    cpu_energy_kwh: float
    ram_energy_kwh: float
    total_emissions_kg: float
    emissions_rate_kg_per_hour: float
    grid_intensity_g_per_kwh: float
    pue_factor: float
    embodied_carbon_kg: float
    carbon_offset_kg: float
    net_emissions_kg: float
    region: str
    timestamp: str
    equivalent_car_days: float
    equivalent_flights_kg: float
    metadata: Dict[str, Any] = field(default_factory=dict)


class GridIntensityProvider:
    """Fetches real-time grid carbon intensity data."""

    PROVIDERS = {
        "electricitymap": "https://api.co2signal.com/v1/latest",
        "watttime": "https://api.watttime.org/v3/latest",
    }

    def __init__(self, api_key: Optional[str] = None, provider: str = "electricitymap"):
        self.api_key = api_key
        self.provider = provider
        self._cache: Dict[str, tuple] = {}
        self._cache_ttl = 300  # 5 minutes

    def get_intensity(
        self, country_code: str, zone_id: Optional[str] = None
    ) -> float:
        """Get current grid carbon intensity in gCO2/kWh."""
        cache_key = f"{country_code}_{zone_id}"
        now = time.time()

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

        if self.provider == "electricitymap":
            intensity = self._fetch_electricitymap(country_code, zone_id)
        elif self.provider == "watttime":
            intensity = self._fetch_watttime(country_code, zone_id)
        else:
            intensity = self._get_default_intensity(country_code)

        self._cache[cache_key] = (now, intensity)
        return intensity

    def _fetch_electricitymap(
        self, country_code: str, zone_id: Optional[str]
    ) -> float:
        """Fetch from ElectricityMap/CO2Signal API."""
        headers = {"auth-token": self.api_key} if self.api_key else {}
        params = {"countryCode": country_code}
        if zone_id:
            params["zone"] = zone_id

        try:
            resp = requests.get(
                self.PROVIDERS["electricitymap"],
                headers=headers,
                params=params,
                timeout=10,
            )
            resp.raise_for_status()
            data = resp.json()
            return data["data"]["carbonIntensity"]
        except Exception as e:
            logger.warning(f"Failed to fetch grid intensity: {e}")
            return self._get_default_intensity(country_code)

    def _fetch_watttime(self, country_code: str, zone_id: Optional[str]) -> float:
        """Fetch from WattTime API."""
        if not self.api_key:
            return self._get_default_intensity(country_code)

        try:
            resp = requests.get(
                self.PROVIDERS["watttime"],
                headers={"Authorization": f"Bearer {self.api_key}"},
                timeout=10,
            )
            resp.raise_for_status()
            data = resp.json()
            return data["data"]["zone"]["carbonIntensity"]
        except Exception as e:
            logger.warning(f"Failed to fetch WattTime intensity: {e}")
            return self._get_default_intensity(country_code)

    @staticmethod
    def _get_default_intensity(country_code: str) -> float:
        """Default grid intensity values (gCO2/kWh) by country."""
        defaults = {
            "USA": 386, "CHN": 555, "IND": 708, "DEU": 338,
            "GBR": 220, "FRA": 56, "CAN": 120, "BRA": 75,
            "JPN": 480, "KOR": 415, "AUS": 530, "NOR": 10,
            "SWE": 13, "ISL": 0, "NZL": 120, "TWN": 500,
        }
        return defaults.get(country_code, 400)


class CarbonTracker:
    """Main carbon tracking system for ML training."""

    def __init__(self, config: CarbonConfig):
        self.config = config
        self.grid_provider = GridIntensityProvider()
        self._start_time: Optional[float] = None
        self._energy_samples: list = []
        self._tracker: Optional[EmissionsTracker] = None

    def start(self):
        """Initialize the carbon tracker."""
        self._start_time = time.time()
        logger.info(
            f"CarbonTracker started for experiment: {self.config.experiment_name}"
        )

        self._tracker = EmissionsTracker(
            project_name=self.config.project_name,
            experiment_id=self.config.experiment_name,
            experiment_name=self.config.experiment_name,
            output_dir="./carbon_logs",
            save_to_file=True,
            country_code=self.config.country_code,
            region=self.config.cloud_region,
            gpu_ids=list(range(self.config.gpu_count)),
        )
        self._tracker.start()

    def stop(self) -> CarbonReport:
        """Stop tracking and generate the final report."""
        if self._tracker:
            total_emissions = self._tracker.stop()
        else:
            total_emissions = 0.0

        duration = time.time() - self._start_time
        duration_hours = duration / 3600

        grid_intensity = self.grid_provider.get_intensity(
            self.config.country_code
        )

        gpu_energy = self._estimate_gpu_energy(duration)
        cpu_energy = self._estimate_cpu_energy(duration)
        ram_energy = self._estimate_ram_energy(duration)
        total_energy = gpu_energy + cpu_energy + ram_energy

        pue_adjusted_energy = total_energy * self.config.pue_factor
        operational_emissions = pue_adjusted_energy * grid_intensity / 1000

        embodied_carbon = self._estimate_embodied_carbon(duration)
        total_emissions_kg = operational_emissions + embodied_carbon

        net_emissions = total_emissions_kg - self.config.carbon_offset_per_kg

        report = CarbonReport(
            experiment_name=self.config.experiment_name,
            model_name=self.config.project_name,
            duration_hours=round(duration_hours, 4),
            total_energy_kwh=round(total_energy, 6),
            gpu_energy_kwh=round(gpu_energy, 6),
            cpu_energy_kwh=round(cpu_energy, 6),
            ram_energy_kwh=round(ram_energy, 6),
            total_emissions_kg=round(total_emissions_kg, 6),
            emissions_rate_kg_per_hour=round(total_emissions_kg / max(duration_hours, 0.001), 6),
            grid_intensity_g_per_kwh=grid_intensity,
            pue_factor=self.config.pue_factor,
            embodied_carbon_kg=round(embodied_carbon, 6),
            carbon_offset_kg=self.config.carbon_offset_per_kg,
            net_emissions_kg=round(net_emissions, 6),
            region=self.config.cloud_region or self.config.country_code,
            timestamp=datetime.now(timezone.utc).isoformat(),
            equivalent_car_days=round(total_emissions_kg / 4.6, 4),
            equivalent_flights_kg=round(total_emissions_kg / 255, 6),
        )

        self._log_to_database(report)
        return report

    def _estimate_gpu_energy(self, duration_seconds: float) -> float:
        """Estimate GPU energy consumption in kWh."""
        gpu_tdp = {
            "NVIDIA A100": 300, "NVIDIA H100": 350,
            "NVIDIA V100": 300, "NVIDIA A10": 150,
            "NVIDIA RTX 4090": 450, "NVIDIA RTX 3090": 350,
        }
        tdp = gpu_tdp.get(self.config.gpu_type, 250)
        avg_utilization = 0.75
        energy_kwh = (
            tdp * avg_utilization * self.config.gpu_count * duration_seconds
        ) / (1000 * 3600)
        return energy_kwh

    def _estimate_cpu_energy(self, duration_seconds: float) -> float:
        """Estimate CPU energy consumption in kWh."""
        cpu_tdp = 200
        avg_utilization = 0.30
        energy_kwh = (cpu_tdp * avg_utilization * duration_seconds) / (1000 * 3600)
        return energy_kwh

    def _estimate_ram_energy(self, duration_seconds: float) -> float:
        """Estimate RAM energy consumption in kWh."""
        ram_gb = 256
        watts_per_gb = 3
        energy_kwh = (ram_gb * watts_per_gb * duration_seconds) / (1000 * 3600)
        return energy_kwh

    def _estimate_embodied_carbon(self, duration_seconds: float) -> float:
        """Estimate amortized embodied carbon in kg."""
        gpu_embodied_kg = 300
        server_embodied_kg = 1500
        expected_lifetime_hours = 43800  # 5 years
        total_embodied = (
            gpu_embodied_kg * self.config.gpu_count
            + server_embodied_kg
        )
        amortization_fraction = duration_seconds / (expected_lifetime_hours * 3600)
        return total_embodied * amortization_fraction

    def _log_to_database(self, report: CarbonReport):
        """Log report to PostgreSQL database."""
        if not self.config.enable_db_logging:
            return
        try:
            import psycopg2
            conn = psycopg2.connect("dbname=carbon_tracker")
            cur = conn.cursor()
            cur.execute(
                """INSERT INTO emissions_log
                (experiment_name, model_name, duration_seconds, energy_kwh,
                 emissions_kg, gpu_power_kwh, cpu_power_kwh, ram_power_kwh,
                 emissions_rate_kg_per_hour, region, grid_intensity, pue_factor,
                 carbon_offset_kg, metadata)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""",
                (
                    report.experiment_name, report.model_name,
                    report.duration_hours * 3600, report.total_energy_kwh,
                    report.total_emissions_kg, report.gpu_energy_kwh,
                    report.cpu_energy_kwh, report.ram_energy_kwh,
                    report.emissions_rate_kg_per_hour, report.region,
                    report.grid_intensity_g_per_kwh, report.pue_factor,
                    report.carbon_offset_kg,
                    json.dumps(report.metadata),
                ),
            )
            conn.commit()
            cur.close()
            conn.close()
        except Exception as e:
            logger.warning(f"Failed to log to database: {e}")

# train_with_tracking.py
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from carbon_calculator import CarbonTracker, CarbonConfig


def create_model():
    return nn.Sequential(
        nn.Linear(784, 512), nn.ReLU(), nn.Dropout(0.2),
        nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.2),
        nn.Linear(256, 128), nn.ReLU(),
        nn.Linear(128, 10),
    )


def create_dataloader():
    X = torch.randn(10000, 784)
    y = torch.randint(0, 10, (10000,))
    return DataLoader(TensorDataset(X, y), batch_size=64, shuffle=True)


def train_model(tracker: CarbonTracker):
    model = create_model()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()
    dataloader = create_dataloader()

    for epoch in range(20):
        model.train()
        total_loss = 0
        for batch_X, batch_y in dataloader:
            optimizer.zero_grad()
            output = model(batch_X)
            loss = criterion(output, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1}/20 - Loss: {avg_loss:.4f}")

    return model


if __name__ == "__main__":
    config = CarbonConfig(
        project_name="sustainable-ml-demo",
        experiment_name="baseline-training-v1",
        team_name="ml-research",
        country_code="USA",
        cloud_region="us-east-1",
        gpu_count=1,
        gpu_type="NVIDIA A100",
        pue_factor=1.15,
    )

    tracker = CarbonTracker(config)
    tracker.start()

    model = train_model(tracker)

    report = tracker.stop()

    print("\n" + "=" * 60)
    print("CARBON FOOTPRINT REPORT")
    print("=" * 60)
    print(f"Experiment: {report.experiment_name}")
    print(f"Duration:   {report.duration_hours:.2f} hours")
    print(f"Energy:     {report.total_energy_kwh:.4f} kWh")
    print(f"  GPU:      {report.gpu_energy_kwh:.4f} kWh")
    print(f"  CPU:      {report.cpu_energy_kwh:.4f} kWh")
    print(f"  RAM:      {report.ram_energy_kwh:.4f} kWh")
    print(f"Emissions:  {report.total_emissions_kg:.4f} kg CO₂")
    print(f"Rate:       {report.emissions_rate_kg_per_hour:.4f} kg CO₂/hr")
    print(f"Grid:       {report.grid_intensity_g_per_kwh} gCO₂/kWh")
    print(f"PUE:        {report.pue_factor}")
    print(f"Embodied:   {report.embodied_carbon_kg:.4f} kg CO₂")
    print(f"Net:        {report.net_emissions_kg:.4f} kg CO₂")
    print(f"Equivalent: {report.equivalent_car_days:.2f} car-days")
    print("=" * 60)

# dashboard.py
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
import psycopg2


def load_emissions_data():
    conn = psycopg2.connect("dbname=carbon_tracker")
    df = pd.read_sql(
        """SELECT experiment_name, timestamp, emissions_kg,
                  energy_kwh, grid_intensity, region
           FROM emissions_log
           ORDER BY timestamp""",
        conn,
    )
    conn.close()
    return df


def create_dashboard(df):
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=(
            "Cumulative Emissions", "Emissions by Experiment",
            "Grid Intensity Over Time", "Energy Breakdown"
        ),
    )

    fig.add_trace(
        go.Scatter(
            x=df["timestamp"],
            y=df["emissions_kg"].cumsum(),
            mode="lines+markers",
            name="Cumulative CO₂ (kg)",
            line=dict(color="#16a34a", width=3),
        ),
        row=1, col=1,
    )

    exp_totals = df.groupby("experiment_name")["emissions_kg"].sum().sort_values()
    fig.add_trace(
        go.Bar(
            x=exp_totals.values,
            y=exp_totals.index,
            orientation="h",
            marker_color="#22c55e",
            name="Total Emissions",
        ),
        row=1, col=2,
    )

    fig.update_layout(
        height=800,
        title_text="ML Training Carbon Footprint Dashboard",
        showlegend=True,
        template="plotly_white",
    )

    fig.write_html("carbon_dashboard.html")
    fig.show()


if __name__ == "__main__":
    df = load_emissions_data()
    create_dashboard(df)