# Copyright (c) 2019-2021 by University of Kassel, Tu Dortmund, RWTH Aachen University and Fraunhofer
# Institute for Energy Economics and Energy System Technology (IEE) Kassel and individual
# contributors (see AUTHORS file for details). All rights reserved.

import pytest
import os
from copy import deepcopy
import numpy as np
import pandas as pd
import pandapower as pp
from pandapower.networks import example_simple

from simbench import sb_dir
from simbench.converter import csv2pp, csv_data2pp, pp2csv, pp2csv_data, \
    convert_parallel_branches, read_csv_data, ensure_full_column_data_existence, \
    avoid_duplicates_in_column, merge_busbar_coordinates

try:
    import pplog as logging
except ImportError:
    import logging

logger = logging.getLogger(__name__)

simbench_converter_test_path = os.path.join(sb_dir, "test", "converter")
test_network_path = os.path.join(simbench_converter_test_path, "test_network")
test_output_folder_path = os.path.join(simbench_converter_test_path, "test_network_output_folder")

__author__ = 'smeinecke'


def test_convert_to_parallel_branches():
    # create test grid
    net = pp.create_empty_network()
    pp.create_bus(net, 110)
    pp.create_buses(net, 2, 20)

    # --- transformers & corresponding switches
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 1")
    pp.create_switch(net, 1, 0, "t", name="Tr-Switch 1")
    # only name changed:
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 2")
    pp.create_switch(net, 1, 1, "t", name="Tr-Switch 2")
    # only max_loading changed:
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 1", max_loading_percent=50)
    pp.create_switch(net, 1, 2, "t", name="Tr-Switch 1")
    # only switch position changed:
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 1")
    pp.create_switch(net, 1, 3, "t", closed=False, name="Tr-Switch 1")
    # only switch missing:
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 1")
    # only name and std_type changed:
    pp.create_transformer(net, 0, 1, "25 MVA 110/20 kV", name="Trafo 3")
    pp.create_switch(net, 1, 5, "t", name="Tr-Switch 3")
    # only name changed and switch added:
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 4")
    pp.create_switch(net, 1, 6, "t", name="Tr-Switch 4a")
    pp.create_switch(net, 0, 6, "t", name="Tr-Switch 4b")
    # only name and parallel changed:
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="Trafo 5", parallel=2)
    pp.create_switch(net, 1, 7, "t", name="Tr-Switch 5")

    # --- lines & corresponding switches
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 1")
    pp.create_switch(net, 2, 0, "l", name="L-Switch 1")
    # only name changed:
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 2")
    pp.create_switch(net, 2, 1, "l", name="L-Switch 2")
    # only max_loading changed:
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 1", max_loading_percent=50)
    pp.create_switch(net, 2, 2, "l", name="L-Switch 1")
    # only switch position changed:
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 1")
    pp.create_switch(net, 2, 3, "l", closed=False, name="L-Switch 1")
    # only switch missing:
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 1")
    # only name and std_type changed:
    pp.create_line(net, 1, 2, 1.11, "48-AL1/8-ST1A 20.0", name="Line 3")
    pp.create_switch(net, 2, 5, "l", name="L-Switch 3")
    # only name changed and switch added:
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 4")
    pp.create_switch(net, 2, 6, "l", name="L-Switch 4a")
    pp.create_switch(net, 1, 6, "l", name="L-Switch 4b")
    # only name and parallel changed:
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="Line 5", parallel=2)
    pp.create_switch(net, 2, 7, "l", name="L-Switch 5")
    # only name and from_bus <-> to_bus changed:
    pp.create_line(net, 2, 1, 1.11, "94-AL1/15-ST1A 20.0", name="Line 6")
    pp.create_switch(net, 2, 8, "l", name="L-Switch 6")

    net1 = deepcopy(net)
    net2 = deepcopy(net)
    net3 = deepcopy(net)

    # complete
    convert_parallel_branches(net1, multiple_entries=False)
    for elm in ["trafo", "line"]:
        assert sorted(net1[elm].index) == [0, 2, 3, 4, 5, 6]
    assert list(net1["trafo"].parallel.values) == [4] + [1]*5
    assert list(net1["line"].parallel.values) == [5] + [1]*5

    # only line
    convert_parallel_branches(net2, multiple_entries=False, elm_to_convert=["line"])
    assert pp.dataframes_equal(net2.line, net1.line)
    assert pp.dataframes_equal(net2.trafo, net.trafo)

    # only exclude "max_loading_percent"
    convert_parallel_branches(net3, multiple_entries=False, exclude_cols_from_parallel_finding=[
                                      "name", "parallel", "max_loading_percent"])
    for elm in ["trafo", "line"]:
        assert sorted(net3[elm].index) == [0, 3, 4, 5, 6]
    assert list(net3["trafo"].parallel.values) == [5] + [1]*4
    assert list(net3["line"].parallel.values) == [6] + [1]*4


def test_convert_parallel_branches():
    # create test grid
    net = pp.create_empty_network()
    pp.create_bus(net, 110)
    pp.create_buses(net, 4, 20)
    pp.create_ext_grid(net, 0)
    pp.create_load(net, 4, 1e3, 4e2)
    pp.create_transformer(net, 0, 1, "40 MVA 110/20 kV", name="sd", parallel=3)
    pp.create_switch(net, 1, 0, "t", name="dfjk")
    pp.create_line(net, 1, 2, 1.11, "94-AL1/15-ST1A 20.0", name="sdh", parallel=2)
    pp.create_switch(net, 2, 0, "l", name="dfsdf")
    pp.create_line(net, 2, 3, 1.11, "94-AL1/15-ST1A 20.0", name="swed", parallel=1)
    pp.create_line(net, 3, 4, 1.11, "94-AL1/15-ST1A 20.0", name="sdhj", parallel=3)
    pp.create_switch(net, 3, 2, "l", name="dfdfg")
    pp.create_switch(net, 4, 2, "l", False, name="dfhgj")
    # check test grid
    assert net.trafo.shape[0] == 1
    assert net.line.shape[0] == 3
    assert net.switch.shape[0] == 4

    convert_parallel_branches(net)
    # test parallelisation
    assert net.trafo.shape[0] == 3
    assert net.line.shape[0] == 6
    assert net.switch.shape[0] == 11

    net1 = deepcopy(net)
    net1.switch.closed.loc[4] = False
    convert_parallel_branches(net, multiple_entries=False)
    convert_parallel_branches(net1, multiple_entries=False)
    # test sum up of parallels
    assert net.trafo.shape[0] == 1
    assert net.line.shape[0] == 3
    assert net.switch.shape[0] == 4
    assert net1.trafo.shape[0] == 1
    assert net1.line.shape[0] == 4
    assert net1.switch.shape[0] == 5


def test_test_network():
    net = csv2pp(test_network_path, no_generic_coord=True)

    # test min/max ratio
    for elm in pp.pp_elements(bus=False, branch_elements=False, other_elements=False):
        if "min_p_mw" in net[elm].columns and "max_p_mw" in net[elm].columns:
            isnull = net[elm][["min_p_mw", "max_p_mw"]].isnull().any(1)
            assert (net[elm].min_p_mw[~isnull] <= net[elm].max_p_mw[~isnull]).all()
        if "min_q_mvar" in net[elm].columns and "max_q_mvar" in net[elm].columns:
            isnull = net[elm][["min_q_mvar", "max_q_mvar"]].isnull().any(1)
            assert (net[elm].min_q_mvar[~isnull] <= net[elm].max_q_mvar[~isnull]).all()

    pp2csv(net, test_output_folder_path, export_pp_std_types=False, drop_inactive_elements=False)

    # --- test equality of exported csv data and given csv data
    csv_orig = read_csv_data(test_network_path, ";")
    csv_exported = read_csv_data(test_output_folder_path, ";")

    all_eq = True
    for tablename in csv_orig.keys():
        try:
            eq = pp.dataframes_equal(csv_orig[tablename], csv_exported[tablename], tol=1e-7)
            if not eq:
                logger.error("csv_orig['%s'] and csv_exported['%s'] differ." % (tablename,
                                                                                tablename))
                logger.error(csv_orig[tablename].head())
                logger.error(csv_exported[tablename].head())
                logger.error(csv_orig[tablename].dtypes)
                logger.error(csv_exported[tablename].dtypes)
        except ValueError:
            eq = False
            logger.error("dataframes_equal did not work for %s." % tablename)
        all_eq &= eq
    assert all_eq


def test_example_simple():
    net = example_simple()

    # --- fix scaling
    net.load["scaling"] = 1.

    # --- add some additional data
    net.bus["subnet"] = ["net%i" % i for i in net.bus.index]
    pp.create_measurement(net, "i", "trafo", np.nan, np.nan, 0, "hv", name="1")
    pp.create_measurement(net, "i", "line", np.nan, np.nan, 1, "to", name="2")
    pp.create_measurement(net, "v", "bus", np.nan, np.nan, 0, name="3")

    net.shunt["max_step"] = np.nan
    stor = pp.create_storage(net, 6, 0.01, 0.1, -0.002, 0.05, 80, name="sda", min_p_mw=-0.01,
                             max_p_mw=0.008, min_q_mvar=-0.01, max_q_mvar=0.005)
    net.storage.loc[stor, "efficiency_percent"] = 90
    net.storage.loc[stor, "self-discharge_percent_per_day"] = 0.3
    pp.create_dcline(net, 4, 6, 0.01, 0.1, 1e-3, 1.0, 1.01, name="df", min_q_from_mvar=-0.01)
    pp.runpp(net)
    to_drop = pp.create_bus(net, 7, "to_drop")

    # --- add names to elements
    for i in pp.pp_elements():
        net[i] = ensure_full_column_data_existence(net, i, 'name')
        avoid_duplicates_in_column(net, i, 'name')

    # --- create geodata
    net.bus_geodata["x"] = [0, 1, 2, 3, 4, 5, 5, 3.63]
    net.bus_geodata["y"] = [0]*5+[-5, 5, 2.33]
    merge_busbar_coordinates(net)

    # --- convert
    csv_data = pp2csv_data(net, export_pp_std_types=True, drop_inactive_elements=True)
    net_from_csv_data = csv_data2pp(csv_data)

    # --- adjust net appearance
    net_from_csv_data.converged = net.converged
    pp.drop_buses(net, [to_drop])
    del net["OPF_converged"]
    net.load["type"] = np.nan
    del net_from_csv_data["substation"]
    del net_from_csv_data["profiles"]
    if "power_station_trafo" in net.gen.columns:
        assert net.gen["power_station_trafo"].isnull().all()
        del net.gen["power_station_trafo"]
    for key in net.keys():
        if isinstance(net[key], pd.DataFrame):
            # drop unequal columns
            dummy_columns = net[key].columns
            extra_columns = net_from_csv_data[key].columns.difference(dummy_columns)
            net_from_csv_data[key].drop(columns=extra_columns, inplace=True)
            # drop result table rows
            if "res_" in key:
                if not key == "res_bus":
                    net[key].drop(net[key].index, inplace=True)
                else:
                    net[key].loc[:, ["p_mw", "q_mvar"]] = np.nan
            # adjust dtypes
            if net[key].shape[0]:
                try:
                    net_from_csv_data[key] = net_from_csv_data[key].astype(dtype=dict(net[
                        key].dtypes))
                except:
                    logger.error("dtype adjustment of %s failed." % key)
    # compare std_types as dataframes -> allow tolerance
    for key, vals in net["std_types"].items():
        net["std_types|%s" % key] = pd.DataFrame(vals).T.apply(pd.to_numeric, errors='ignore')
    del net["std_types"]
    for key, vals in net_from_csv_data["std_types"].items():
        net_from_csv_data["std_types|%s" % key] = pd.DataFrame(vals).T.apply(
            pd.to_numeric, errors='ignore')
    del net_from_csv_data["std_types"]

    eq = pp.nets_equal(net, net_from_csv_data, tol=1e-7)
    assert eq


if __name__ == "__main__":
    if 0:
        pytest.main([__file__, "-xs"])
    else:
        test_convert_to_parallel_branches()
        test_convert_parallel_branches()
        test_test_network()
        test_example_simple()
        pass
