# -*- coding: utf-8 -*-
# SPDX-FileCopyrightText: : 2020-2023 The PyPSA-Eur Authors
#
# SPDX-License-Identifier: MIT
"""
Build district heat shares at each node, depending on investment year.
"""

import pandas as pd

from prepare_sector_network import get

import logging


logger = logging.getLogger(__name__)


if __name__ == "__main__":
    if "snakemake" not in globals():
        from _helpers import mock_snakemake

        snakemake = mock_snakemake(
            "build_heat_demands",
            simpl="",
            clusters=48,
        )

    investment_year = int(snakemake.wildcards.planning_horizons[-4:])

    pop_layout = pd.read_csv(snakemake.input.clustered_pop_layout,
                             index_col=0)

    district_heat_share = pd.read_csv(snakemake.input.district_heat_share,
                                      index_col=0).squeeze()

    # make ct-based share nodal
    district_heat_share = district_heat_share.loc[pop_layout.ct]
    district_heat_share.index = pop_layout.index

    # total urban population per country
    ct_urban = pop_layout.urban.groupby(pop_layout.ct).sum()

    # distribution of urban population within a country
    pop_layout["urban_ct_fraction"] = pop_layout.urban / pop_layout.ct.map(ct_urban.get)

    # fraction of node that is urban
    urban_fraction = pop_layout.urban / pop_layout[["rural", "urban"]].sum(axis=1)

    # maximum potential of urban demand covered by district heating
    central_fraction = snakemake.config["sector"]["district_heating"]["potential"]

    # district heating share at each node
    dist_fraction_node = (
        district_heat_share * pop_layout["urban_ct_fraction"] / pop_layout["fraction"]
    )

    # if district heating share larger than urban fraction -> set urban
    # fraction to district heating share
    urban_fraction = pd.concat([urban_fraction, dist_fraction_node], axis=1).max(axis=1)

    # difference of max potential and today's share of district heating
    diff = (urban_fraction * central_fraction) - dist_fraction_node
    progress = get(snakemake.config["sector"]["district_heating"]["progress"], investment_year)
    dist_fraction_node += diff * progress
    logger.info(
        f"Increase district heating share by a progress factor of {progress:.2%} "
        f"resulting in new average share of {dist_fraction_node.mean():.2%}"
    )

    df = pd.DataFrame(dtype=float)

    df["original district heat share"] = district_heat_share
    df["district fraction of node"] = dist_fraction_node
    df["urban fraction"] = urban_fraction

    df.to_csv(snakemake.output.district_heat_share)