from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt

PROJ_DIR = Path(__file__).parent
CITYGML = 'Grombühl_v2'
URBI_HEAT_DEMAND = 'totalPrimaryEnergyDemandHeating [kWh/a]'
S_DEMAND = 'Simstadt heat demand [MWh / a]'  # Inkl Warmwasser?
U_DEMAND = 'Urbi+ heat demand [MWh / a]'
AREA = 'Area [m²]'
TYPE = 'Usage Type'

# Feldberg : 2.493 (Average : 3.0°C, Min : -18°C)
# Original : 2.487 (Average : 7.9°C, Min : -12°C), Würzburg
# Bensheim : 2.475 (Average : 10.2°C, Min : -10°C)


def get_simstadt_heat_demand():
    heating_csvs = list(PROJ_DIR.glob(f'**/{CITYGML}_DIN18599_HEATING.csv'))
    assert len(heating_csvs) == 1
    hourly_demand = heating_csvs[0]

    simstadt_df = pd.read_csv(hourly_demand, sep=';',
                              skiprows=list(range(19)) + [20],
                              decimal='.')

    simstadt_df = simstadt_df.rename({
        'Total Yearly Heat+DHW demand': S_DEMAND,
        'PrimaryUsageZoneArea': AREA,
        'PrimaryUsageZoneType': TYPE,
    }, axis='columns')

    simstadt_df[S_DEMAND] = simstadt_df[S_DEMAND] / 1000  # kWh -> MWh
    simstadt_df = simstadt_df.set_index('GMLId')

    by_type_df = simstadt_df.groupby([TYPE]).sum(
        numeric_only=True)[[AREA, S_DEMAND]]

    by_type_df['Specific heat demand [kWh / (m².a)]'] = by_type_df[S_DEMAND] / \
        by_type_df[AREA] * 1000

    print(by_type_df.round(1).to_string())

    return simstadt_df


def get_urbi_heat_demand():
    urbi_csvs = list(PROJ_DIR.glob(f'**/{CITYGML}.csv'))
    assert len(urbi_csvs) == 1
    urbi_csv_filename = urbi_csvs[0]
    urbi_results = {}
    with open(urbi_csv_filename) as csv_file:
        cells = next(csv_file).strip().split(';')
        heat_demand_column = cells.index(URBI_HEAT_DEMAND)
        id_column = cells.index('ID')
        for line in csv_file:
            cells = line.strip().split(';')
            urbi_results[cells[id_column]] = round(
                float(cells[heat_demand_column]), 0) / 1000  # MWh

    urbi_df = pd.DataFrame.from_dict(
        urbi_results, orient='index', columns=[U_DEMAND])
    return urbi_df


def compare_results(simstadt_df, urbi_df):
    print(urbi_df)
    simstadt_urbi = simstadt_df.join(urbi_df)[[S_DEMAND, U_DEMAND]]
    simstadt_urbi = simstadt_urbi.dropna()

    simstadt_urbi.round(1).to_csv(
        f'simstadt_vs_urbi_heat_demands_{CITYGML}.csv', sep=';', decimal='.')

    plt.rcParams["figure.figsize"] = (10, 10)
    fig = simstadt_urbi.plot.scatter(
        x=S_DEMAND, y=U_DEMAND, title=f'Urbi vs SimStadt Heat-demand\n{CITYGML}')
    max_heat_demand = simstadt_urbi.max().max()
    fig.plot([0, max_heat_demand], [0, max_heat_demand])
    fig.axis('equal')
    plt.savefig(
        f'simstadt_vs_urbi_heat_demands_{CITYGML}.png', bbox_inches='tight')

    print("Correlation:")
    print(simstadt_urbi.corr())

    print()
    print(simstadt_urbi.round(1))

    print()
    print("Average Urbi / SimStadt heat demand ratio:")
    print((simstadt_urbi[U_DEMAND] / simstadt_urbi[S_DEMAND]).mean())


if __name__ == '__main__':
    compare_results(get_simstadt_heat_demand(), get_urbi_heat_demand())
