# -*- coding: utf-8 -*-
"""
/***************************************************************************
 GenerateSwmmInp
                                 A QGIS plugin
 This plugin generates SWMM Input files
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2021-07-09
        copyright            : (C) 2023 by Jannik Schilling
        email                : jannik.schilling@posteo.de
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

__author__ = 'Jannik Schilling'
__date__ = '2023-05-09'
__copyright__ = '(C) 2023 by Jannik Schilling'


import pandas as pd
import numpy as np
from .g_s_various_functions import check_columns
from .g_s_defaults import (
    def_infiltration_types,
    def_qgis_fields_dict,
    def_sections_dict
)


def get_subcatchments_from_layer(subcatchments_df, main_infiltration_method):
    """
    reads subcatchment shapefile
    """

    def rename_for_infiltation(subc_row, main_infiltration_method):
        """
        selects and renames different columns
        according to the columns 'InfMethod'
        """
        infiltr_method = subc_row['InfMethod']
        if pd.isna(infiltr_method):
            # take main infiltration method, if not given for subcatchment
            infiltr_method = main_infiltration_method
        if infiltr_method in ['GREEN_AMPT', 'MODIFIED_GREEN_AMPT']:
            subc_row = subc_row.rename({
                'SuctHead': 'Param1',
                'Conductiv': 'Param2',
                'InitDef': 'Param3'
            })
            subc_row['Param4'] = ''
            subc_row['Param5'] = ''
        if infiltr_method in ['HORTON', 'MODIFIED_HORTON']:
            subc_row = subc_row.rename({
                'MaxRate': 'Param1',
                'MinRate': 'Param2',
                'Decay': 'Param3',
                'DryTime': 'Param4',
                'MaxInf': 'Param5'
            })
        if infiltr_method == 'CURVE_NUMBER':
            subc_row = subc_row.rename({
                'CurveNum': 'Param1',
                'Conductiv': 'Param2',
                'DryTime': 'Param3'
            })
            subc_row['Param4'] = ''
            subc_row['Param5'] = ''
        if infiltr_method is None:
            subc_row['Param1'] = ''
            subc_row['Param2'] = ''
            subc_row['Param3'] = ''
            subc_row['Param4'] = ''
            subc_row['Param5'] = ''
        return subc_row
    subcatchments_df = subcatchments_df.apply(
        rename_for_infiltation,
        axis=1,
        args=(main_infiltration_method, )
    )
    subcatchments_df['CurbLen'] = subcatchments_df['CurbLen'].fillna('0')
    subcatchments_df['SnowPack'] = subcatchments_df['SnowPack'].fillna('')
    subcatchments_df['PctRouted'] = subcatchments_df['PctRouted'].fillna(100)
    subcatchments_df = subcatchments_df.reset_index(drop=True)
    # select columns
    infiltration_df = subcatchments_df[def_sections_dict['INFILTRATION']]
    subareas_df = subcatchments_df[def_sections_dict['SUBAREAS']]
    subcatchments_df = subcatchments_df[def_sections_dict['SUBCATCHMENTS']]
    return subcatchments_df, subareas_df, infiltration_df

def adjust_infiltration_inp_lines(
    inp_line,
    main_infiltration_method
):
    """
    adjusts the line length in the infiltration section
    """
    if inp_line[-1] in def_infiltration_types:
        current_infiltration_method = inp_line[-1]
        inp_line = inp_line[:-1]
    else:
        current_infiltration_method = main_infiltration_method
    if len(inp_line) == 4:
        # fill non-HORTON or missing values
        inp_line = inp_line+[np.nan, np.nan]
    if len(inp_line) == 5:
        # fill missing values
        inp_line = inp_line+[np.nan]
    inp_line = inp_line+[current_infiltration_method]
    return inp_line

def create_subcatchm_attributes_from_inp_df(
    all_subcatchments,
    all_subareas,
    all_infiltr
):
    """
    creates pd.Dataframes from a pd.DataFrame of subcatchment attributes (from an inp file)
    :param pd.DataFrame all_subcatchments
    :param pd.DataFrame all_subareas
    :param pd.DataFrame all_infiltr
    """
    InfiltrDtypes = [
        'InfMethod',
        'SuctHead',
        'Conductiv',
        'InitDef',
        'MaxRate',
        'MinRate',
        'Decay',
        'DryTime',
        'MaxInf',
        'CurveNum'
    ]

    def create_infiltr_df(infiltr_row):
        """
        creates a pd.DataFrame for infiltration values
        :param pd.Series infiltr_row
        """
        if infiltr_row['InfMethod'] in ['GREEN_AMPT', 'MODIFIED_GREEN_AMPT']:
            infiltr_row = infiltr_row.rename({
                'Param1': 'SuctHead',
                'Param2': 'Conductiv',
                'Param3': 'InitDef'
            })
            infiltr_row = infiltr_row.drop(['Param4', 'Param5'])
            cols_not_in_infilt = [k for k in InfiltrDtypes if k not in infiltr_row.index]  # missing columns
            for c in cols_not_in_infilt:
                infiltr_row[c] = np.nan
        if infiltr_row['InfMethod'] in ['HORTON', 'MODIFIED_HORTON']:
            infiltr_row = infiltr_row.rename({
                'Param1': 'MaxRate',
                'Param2': 'MinRate',
                'Param3': 'Decay',
                'Param4': 'DryTime',
                'Param5': 'MaxInf'
            })
            cols_not_in_infilt = [k for k in InfiltrDtypes if k not in infiltr_row.index]  # missing columns
            for c in cols_not_in_infilt:
                infiltr_row[c] = np.nan
        if infiltr_row['InfMethod'] == 'CURVE_NUMBER':
            infiltr_row = infiltr_row.rename({
                'Param1': 'CurveNum',
                'Param2': 'Conductiv',
                'Param3': 'DryTime'
            })
            infiltr_row = infiltr_row.drop(['Param4', 'Param5'])
            cols_not_in_infilt = [k for k in InfiltrDtypes if k not in infiltr_row.index]  # missing columns
            for c in cols_not_in_infilt:
                infiltr_row[c] = np.nan
        return infiltr_row
    all_infiltr = all_infiltr.apply(lambda x: create_infiltr_df(x), axis=1)
    all_infiltr = all_infiltr[['Name']+InfiltrDtypes]
    all_subcatchments = all_subcatchments.join(all_subareas.set_index('Name'), on='Name')
    all_subcatchments = all_subcatchments.join(all_infiltr.set_index('Name'), on='Name')
    return all_subcatchments


# for raingages
def get_raingage_list_from_inp(rg_line):
    """
    creates a list of raingage values in the correct order from an inp line
    :param list rg_line
    :return list
    """
    if rg_line[4] == 'TIMESERIES':
        rg_source = {
            'DataSource': 'TIMESERIES',
            'SeriesName': rg_line[5],
            'FileName': np.nan,
            'StationID': np.nan,
            'RainUnits': np.nan
        }
    else:  # FILE
        rg_source = {
            'DataSource': 'FILE',
            'SeriesName': np.nan,
            'FileName': rg_line[5],
            'StationID': rg_line[6],
            'RainUnits': rg_line[7]
        }
    interval_split = str(rg_line[2]).split(':')  # Interval splitted in HH:mm
    if len(interval_split) == 1:
        interval_split = interval_split + ['00']  # if only hours
    if len(interval_split[0]) == 1:
        interval_split[0] = '0' + interval_split[0]  # if one digit hour
    interval = interval_split[0] + ':' + interval_split[1]
    rg_list = [
        rg_line[0],  # Name
        rg_line[1],  # Format
        interval,
        rg_line[3]  # SCF
    ] + list(rg_source.values())
    return rg_list


def get_raingage_from_qgis_row(rg_row):
    """
    adjusts columns in a row from a QGIS raingage layer
    :param pd.series rg_row
    """
    if rg_row['DataSource'] == 'TIMESERIES':
        rg_row['SourceName'] = rg_row['SeriesName']
        rg_row['StationID'] = ''
        rg_row['RainUnits'] = ''
    else:  # FILE
        rg_row['SourceName'] = rg_row['FileName']
    # unused columns (FileName and SeriesName) are deleted later on
    return rg_row