import io
import os
import typing
from collections import OrderedDict
from pathlib import Path

import numpy as np
try:
    import pandas as pd
except ImportError:
    from ..stubs import pandas as pd

from .weir import Weir
from ..output import Output, OutputCollection


if typing.TYPE_CHECKING:
    from ..parsers.units.handler import Handler


class QHControl(Weir):

    @staticmethod
    def complete_unit_type_name() -> str:
        return 'QH CONTROL_'

    def has_qh(self) -> bool:
        return True

    def map_nwk_attributes(self, field_map: dict, unit: 'Handler') -> OrderedDict:
        d = super().map_nwk_attributes(field_map, unit)
        d['Type'] = 'M'
        d['DS_Invert'] = None
        return d

    def get_qh(self) -> Output:
        self.qh = super().get_qh()
        df = self.calc_qh_matrix()
        buf = io.StringIO()
        buf.write(f'! Generated by fm_to_estry. Source: {self.dat.name}/{self.unit.uid}\n')
        df.to_csv(buf, lineterminator='\n', float_format='%.3f')
        self.qh.content = buf.getvalue()
        return self.qh

    def calc_qh_matrix(self) -> pd.DataFrame:
        # convert level to depth
        curve = self.unit.qh.rename(columns={'Elevation': 'd', 'Flow': 'q'})
        curve['d'] = curve['d']- self.unit.zc
        # empty matrix
        matrix = self.empty_matrix(curve)
        # populate matrix
        for i, d2 in enumerate(matrix.index):
            for j, d1 in enumerate(matrix.columns):
                matrix.iloc[i, j] = self.extract_flow(d1, d2, curve)
        return matrix

    def empty_matrix(self, curve: pd.DataFrame) -> pd.DataFrame:
        npts = 11 - curve.shape[0] + 1
        inc = curve['d'].iloc[-1] / npts
        depths = [inc * x for x in range(npts + 1)]
        for dep in curve['d']:
            if dep not in depths:
                depths.append(dep)
        depths = sorted(depths)
        df = pd.DataFrame(np.zeros((len(depths), len(depths))), columns=depths, index=depths)
        df.index.name = 1.
        return df

    def extract_flow(self, d1: float, d2: float, curve: pd.DataFrame) -> float:
        sign = -1. if d2 > d1 else 1.
        if d1 < d2:
            d1, d2 = d2, d1
        return self.calc_flow(d1, d2, curve) * sign

    def calc_flow(self, d1: float, d2: float, curve: pd.DataFrame) -> float:
        if np.isclose(d1, 0.) or np.isclose(d1, d2):
            return 0.
        if d2 / d1 < self.unit.m:  # free-flow
            drownf = 1.
        else:
            drownf = np.sqrt(((1 - d2) / d1) / (1 - self.unit.m))
            if drownf < 0.3:
                drownf = ((1 - d2) / d1) / (0.3 * (1 - self.unit.m))
        return np.interp(d1, curve['d'], curve['q']) * drownf
