import sys, math
import numpy as np
import pandas as pd
import pycountry
import pycountry_convert as pc
import geonamescache
D = sys.argv[1] if len(sys.argv) > 1 else "."
EXPOSANT, PLANCHER, PLAFOND, NIVEAU_REF = 2.0, 0.30, 2.5, 0.75
BASE = dict(conv=200.0, continu=0.30, alternatif=0.60, poste=50.0, poste_existant=20.0, renforcement=100.0, mer=1.61)
PROJETS = {
    "CHN": ("Changji-Guquan, 12 GW, 3 293 km aériens, 2019, 40,7 Md CNY (5,27 Md€)", 12.0, 3293.0, 0.0, 5270.0),
    "BRA": ("Xingu-Rio, 4 GW, 2 539 km aériens, 2019, 2,14 Md$ (1,90 Md€)", 4.0, 2539.0, 0.0, 1900.0),
    "CHL": ("Kimal-Lo Aguirre, 3 GW, 1 342 km aériens, en chantier, 1,48 Md$ (1,37 Md€)", 3.0, 1342.0, 0.0, 1370.0),
    "USA": ("TransWest Express, 3 GW, 1 172 km aériens, en chantier, 3 Md$ (2,78 Md€)", 3.0, 1172.0, 0.0, 2780.0),
}
SOUTERRAIN = {"DEU": ("SuedLink, 4 GW, 700 km souterrains imposés par la loi, en chantier, 10 Md€", 4.0, 700.0, 10000.0)}
VERIF = ("Viking Link, Royaume-Uni et Danemark, 1,4 GW, 765 km dont 620 sous la mer, 2023, 1,7 Md£ (1,98 Md€)", 1.4, 145.0, 620.0, 1980.0)
def facteur_projet(gw, terre, mer, meur):
    par_gw = meur / gw
    fixe = BASE["conv"] * 0.6 + BASE["mer"] * 0.6 * mer
    var = BASE["conv"] * 0.4 + BASE["continu"] * terre + BASE["mer"] * 0.4 * mer
    return (par_gw - fixe) / var
def continent(iso3):
    try:
        a2 = pycountry.countries.get(alpha_3=iso3).alpha_2 if iso3 != "XKX" else "XK"
        return pc.country_alpha2_to_continent_code(a2)
    except Exception:
        return "AS" if iso3 == "TWN" else "EU" if iso3 == "OWID_KOS" else "?"
w = pd.read_csv(f"{D}/niveau_prix_banque_mondiale.csv")
w.loc[w.iso == "XKX", "iso"] = "OWID_KOS"
w["fiable"] = ((w.annee_ppa - w.annee_change).abs() <= 1) & (w.niveau_prix_usa_1 >= 0.15) & (w.niveau_prix_usa_1 <= 1.6)
w["continent"] = w.iso.map(continent)
med = w[w.fiable].groupby("continent").niveau_prix_usa_1.median()
p = pd.read_csv(f"{D}/electricite_pays_prix_part_fossile.csv")[["iso", "pays"]]
p = p.merge(w[["iso", "niveau_prix_usa_1", "fiable", "annee_ppa"]], on="iso", how="left")
p["continent"] = p.iso.map(continent)
p["source_niveau"] = np.where(p.fiable == True, "Banque mondiale " + p.annee_ppa.fillna(0).astype(int).astype(str), "médiane du continent")
p["niveau_prix"] = np.where(p.fiable == True, p.niveau_prix_usa_1, p.continent.map(med))
p["niveau_prix"] = p.niveau_prix.fillna(med.median())
p["facteur"] = np.clip((p.niveau_prix / NIVEAU_REF) ** EXPOSANT, PLANCHER, PLAFOND)
p["source_facteur"] = "niveau des prix"
for iso, (nom, gw, terre, mer, meur) in PROJETS.items():
    f = facteur_projet(gw, terre, mer, meur)
    p.loc[p.iso == iso, "facteur"] = f
    p.loc[p.iso == iso, "source_facteur"] = "projet réel : " + nom
p["facteur_continu_terre"] = p.facteur
for iso, (nom, gw, terre, meur) in SOUTERRAIN.items():
    p.loc[p.iso == iso, "facteur_continu_terre"] = facteur_projet(gw, terre, 0.0, meur)
    p.loc[p.iso == iso, "source_facteur"] = "niveau des prix ; continu souterrain : " + nom
p["convertisseurs_meur_gw"] = BASE["conv"] * (0.6 + 0.4 * p.facteur)
p["continu_terre_meur_km_gw"] = BASE["continu"] * p.facteur_continu_terre
p["alternatif_400kv_meur_km_gw"] = BASE["alternatif"] * p.facteur
p["poste_neuf_meur_gw"] = BASE["poste"] * p.facteur
p["poste_existant_meur_gw"] = BASE["poste_existant"] * p.facteur
p["renforcement_meur_gw"] = BASE["renforcement"] * p.facteur
p["sous_marin_meur_km_gw"] = BASE["mer"] * (0.6 + 0.4 * p.facteur)
p["ligne_1000km_meur_gw"] = p.convertisseurs_meur_gw + 1000 * p.continu_terre_meur_km_gw
area = {c["iso3"]: float(c["areakm2"]) for c in geonamescache.GeonamesCache().get_countries().values()}
area["OWID_KOS"] = 10887.0
p["rayon_km"] = [round(math.sqrt(area.get(i, 10000.0)) / 2) for i in p.iso]
cols = ["iso", "pays", "continent", "niveau_prix", "source_niveau", "facteur", "source_facteur", "convertisseurs_meur_gw", "continu_terre_meur_km_gw", "alternatif_400kv_meur_km_gw", "poste_neuf_meur_gw", "poste_existant_meur_gw", "renforcement_meur_gw", "sous_marin_meur_km_gw", "ligne_1000km_meur_gw", "rayon_km"]
o = p[cols].copy()
for c in cols[3:4] + cols[5:6] + cols[7:15]:
    o[c] = o[c].astype(float).round(3)
o.to_csv(f"{D}/prix_lignes_electriques_pays.csv", index=False, lineterminator="\n")
nom, gw, terre, mer, meur = VERIF
fg = o.loc[o.iso == "GBR", "facteur"].iloc[0]
fd = o.loc[o.iso == "DNK", "facteur"].iloc[0]
f = (fg + fd) / 2
mod = BASE["conv"] * (0.6 + 0.4 * f) + BASE["continu"] * f * terre + BASE["mer"] * (0.6 + 0.4 * f) * mer
print(f"Viking Link : réel {meur / gw:.0f} M€/GW, modèle {mod:.0f} M€/GW, facteur moyen {f:.2f}")
for iso, (nom, gw, terre, mer, meur) in PROJETS.items():
    fr = o.loc[o.iso == iso, "facteur"].iloc[0]
    nv = p.loc[p.iso == iso, "niveau_prix"].iloc[0]
    print(iso, "facteur projet", round(fr, 3), "formule", round(min(max((nv / NIVEAU_REF) ** EXPOSANT, PLANCHER), PLAFOND), 3))
print(len(o), "pays ;", (p.source_niveau == "médiane du continent").sum(), "à la médiane du continent")
print(med.round(3).to_dict())