from __future__ import annotations

from datetime import datetime
from datetime import timedelta
from types import NoneType  # type: ignore
from typing import get_type_hints
from typing import Iterator

import numpy as np
import pandas as pd


def dftoclass(name: str, df: pd.DataFrame, kind: str = "dataclass") -> None:
    ret = []
    inp = "import numpy as np"
    idt = "from datetime import datetime"
    itd = "from datetime import timedelta"
    iany = "from typing import Any"
    if kind == "dict":
        imports = {"from typing import TypedDict"}
    elif kind == "tuple":
        imports = {"from typing import NamedTuple"}
    else:
        imports = {"from dataclasses import dataclass"}
    for cname, dtype in df.dtypes.to_dict().items():
        k = dtype.kind
        if k == "f":
            ret.append((cname, "float"))
        elif k in {"i", "u"}:
            ret.append((cname, "int"))
        elif k == "c":
            ret.append((cname, "complex"))
        elif k == "b":
            ret.append((cname, "bool"))
        elif k == "S":
            ret.append((cname, "str"))
        elif k == "M":
            ret.append((cname, "datetime"))
            imports.add(idt)
        elif k == "m":
            ret.append((cname, "timedelta"))
            imports.add(itd)
        elif k == "O":
            if len(df) == 0:
                raise RuntimeError(
                    "can't determine type of object in zero row dataframe",
                )
            f = df[cname].iloc[0]
            tf = type(f)
            if tf == str:
                ret.append((cname, "str"))
            elif tf == np.ndarray:
                if len(f) > 0:
                    t = type(f[0])
                    itemt = f"  # item: {t.__name__}"
                else:
                    itemt = ""
                ret.append((cname, f"np.ndarray{itemt}"))
                imports.add(inp)
            else:
                ret.append((cname, "Any"))
                imports.add(iany)
        else:
            ret.append((cname, "Any"))
            imports.add(iany)

    attr = "\n".join(f"    {n}: {t}" for n, t in ret)
    comment = "\n# generated by dftoclass\n\n"
    KINDS = {"dict": "TypedDict", "tuple": "NamedTuple"}
    if kind in KINDS:
        typ = KINDS[kind]
        cls = [
            *imports,
            comment,
            f"class {name}({typ}):",
            attr,
        ]
    else:
        cls = [
            *imports,
            comment,
            "@dataclass",
            f"class {name}:",
            attr,
        ]
    print("\n".join(cls))


KINDS = {
    "f": float,
    "i": int,
    "u": int,
    "c": complex,
    "b": bool,
    "S": str,
    "M": datetime,
    "m": timedelta,
    "O": object,
}


def check_df_columns(df: pd.DataFrame, T: type, full: bool = True) -> set[str]:
    ret = []
    ann = get_type_hints(T)
    missing = set(ann) - set(df.columns)
    if missing:
        ret.append(f"missing columns: {missing}")

    if not full:
        return set(ret)

    n = len(df)
    for cname, dtype in df.dtypes.to_dict().items():
        if cname not in ann:  # don't care about extra columns
            continue
        k = dtype.kind
        typ = ann[cname]
        if not typ == KINDS.get(k):
            if k == "O":  # can't check
                if n == 0:
                    continue
                rtyp = type(df[cname].iloc[0])
                if rtyp == typ:
                    continue
            ret.append(f"{cname}: found type {dtype}, expecting type {typ}")
    return set(ret)


def check_pepxml_columns(df: pd.DataFrame, full: bool = True) -> set[str]:
    from .pepxmltypes import PepXMLRow

    return check_df_columns(df, PepXMLRow, full=full)


def diff(df1: pd.DataFrame, df2: pd.DataFrame, eps: float = 1e-6) -> Iterator[str]:
    if not (len(df1) == len(df2) and (df1.index == df2.index).all()):
        yield "indexes differ"
    if not (
        len(df1.columns) == len(df2.columns) and (df1.columns == df2.columns).all()
    ):
        dc = list(set(df1.columns) ^ set(df2.columns))
        yield f"columns differ: {dc}"

    cols = list(set(df1.columns) & set(df2.columns))
    rsuffix = "_o"
    df = df1[cols].join(df2[cols], rsuffix=rsuffix)
    if len(df) == 0:
        yield "can't join"
        return
    yield from diff_joined(df, rsuffix=rsuffix, eps=eps)


def diff_joined(  # noqa: C901
    df1: pd.DataFrame,
    *,
    rsuffix: str,
    eps: float = 1e-6,
) -> Iterator[str]:
    def dictok(d1, d2) -> bool:
        if not isinstance(d1, dict) and not isinstance(d2, dict):
            ret = d1 == d2
            if isinstance(ret, bool):
                return ret
            return ret.all()

        s = set(d1.keys()) == set(d2.keys())
        if not s:
            return False
        for k in d1:
            ok = dictok(d1[k], d2[k])
            if not ok:
                return False
        return True

    def ndarray_failed(arr: pd.DataFrame, col: str) -> bool:
        a1 = arr[col]
        a2 = arr[col + rsuffix]
        if a1 is None:
            return a2 is not None

        if a2 is None:
            return a1 is not None

        if a1.shape != a2.shape:
            return True
        if a1.dtype != a2.dtype:
            return True

        if len(a1) == 0:
            return False
        if a1.dtype.kind == "f":
            na = np.isnan(a1)
            if (na != np.isnan(a2)).sum() != 0:
                return True
            d = a1[~na] != a2[~na]
        else:
            d = a1 != a2  # ndarray of say str

        if isinstance(d, bool):
            for a, b in zip(a1, a2):
                ok = dictok(a, b)
                if not ok:
                    return True
            return False
        failed = d.sum() != 0

        return failed

    def check_array(s1: pd.Series, s2: pd.Series) -> int:
        t = s1.dtype
        if t != s2.dtype and t.kind != s2.dtype.kind:
            return len(s1)
        na = s1.isna()
        d = na != s2.isna()
        failed = d.sum()
        if not failed:
            if t.kind in {"f", "c", "m", "M"}:
                d = (s1[~na] - s2[~na]).abs() > eps
            else:
                d = s1[~na] != s1[~na]
            failed = d.sum()
        return failed

    def find_type(s: pd.Series) -> type | None:
        for i in range(len(s)):
            rtyp = type(s.iloc[i])
            if rtyp != NoneType:
                return rtyp
        return None

    cols = [c for c in df1.columns if not c.endswith(rsuffix)]
    rcols = [c[: -len(rsuffix)] for c in df1.columns if c.endswith(rsuffix)]

    acols = set(cols) & set(rcols)
    if not acols == set(cols):
        yield "columns differ"
    for col in acols:
        dc1 = df1[col]
        dc2 = df1[col + rsuffix]
        t = dc1.dtype

        if t != dc2.dtype:
            samekind = t.kind == dc2.dtype.kind
            yield f"{col}: types differ {t} != {dc2.dtype} {'continue checking' if samekind else 'skipping'}"
            if not samekind:
                continue

        if t.kind in {"b", "i", "u", "S", "U"}:
            failed = (dc1 != dc2).sum()

        elif t.kind in {"f", "c", "m", "M"}:
            failed = check_array(dc1, dc2)

        elif t.kind == "O":
            rtyp = find_type(dc1)

            if rtyp == np.ndarray:
                failed = (
                    pd.concat([dc1, dc2], axis=1)
                    .apply(
                        ndarray_failed,
                        args=(col,),
                        axis=1,
                    )
                    .sum()
                )

            elif rtyp == str:
                failed = check_array(dc1, dc2)

            else:
                yield f"unknown type: {col}[{rtyp or t}]"
                continue
        else:
            # kind == 'V' void
            yield f"unknown type: {col}[{t}]"
            continue

        if failed:
            yield f"{col}[{t}]: {failed}"
