# SPDX-FileCopyrightText: 2022 Contributors to the Power Grid Model project <dynamic.grid.calculation@alliander.com>
#
# SPDX-License-Identifier: MPL-2.0

from pathlib import Path
from typing import Dict
from unittest.mock import MagicMock, call, mock_open, patch

import numpy as np
import pandas as pd
import pytest
from pytest import param
from structlog.testing import capture_logs

from power_grid_model_io.data_stores.excel_file_store import ExcelFileStore
from power_grid_model_io.data_types.tabular_data import TabularData

from ...utils import MockExcelFile, assert_log_exists

PandasExcelData = Dict[str, pd.DataFrame]


@pytest.fixture()
def objects_excel() -> PandasExcelData:
    return {
        "Nodes": pd.DataFrame([["A", 111.0], ["B", 222.0]], columns=("NAME", "U_NOM")),
        "Lines": pd.DataFrame([["C", "A", "B", "X"]], columns=("NAME", "NODE0", "NODE1", "TYPE")),
    }


@pytest.fixture()
def specs_excel() -> PandasExcelData:
    return {
        "Colors": pd.DataFrame([["A", "Red"], ["B", "Yellow"]], columns=("NODE", "COLOR")),
        "Lines": pd.DataFrame([["X", 123.0]], columns=("LINE", "R")),
    }


def noop(data: pd.DataFrame, *_args, **_kwargs) -> pd.DataFrame:
    return data


def test_constructor():
    # Arrange / Act
    fs = ExcelFileStore()  # no exception

    # Assert
    assert fs.files() == {}


def test_constructor__arg():
    # Arrange / Act
    fs = ExcelFileStore(Path("A.xlsx"))

    # Assert
    assert fs.files() == {"": Path("A.xlsx")}


def test_constructor__arg_kwargs():
    # Arrange / Act
    fs = ExcelFileStore(Path("A.xlsx"), foo=Path("B.xlsx"), bar=Path("C.xls"))

    # Assert
    assert fs.files() == {"": Path("A.xlsx"), "foo": Path("B.xlsx"), "bar": Path("C.xls")}


def test_constructor__kwargs():
    # Arrange / Act
    fs = ExcelFileStore(foo=Path("A.xlsx"), bar=Path("B.xls"))

    # Assert
    assert fs.files() == {"foo": Path("A.xlsx"), "bar": Path("B.xls")}


def test_constructor__too_many_args():
    # Too many (> 1) unnamed arguments
    with pytest.raises(TypeError, match=r"1 to 2.*positional arguments.*3.*given"):
        ExcelFileStore(Path("A.xlsx"), Path("B.xls"))  # type: ignore


def test_constructor__invalid_main_file():
    with pytest.raises(ValueError, match=r"Excel.*\.docx"):
        ExcelFileStore(Path("A.docx"))


def test_constructor__invalid_named_file():
    with pytest.raises(ValueError, match=r"Extra.*\.docx"):
        ExcelFileStore(Path("A.xlsx"), extra=Path("B.docx"))


def test_files__read_only():
    # Arrange
    fs = ExcelFileStore(Path("A.xlsx"), extra=Path("B.xlsx"))
    files = fs.files()

    # Act
    files["extra"] = Path("C.xlsx")

    # Assert
    assert files == {"": Path("A.xlsx"), "extra": Path("C.xlsx")}
    assert fs.files() == {"": Path("A.xlsx"), "extra": Path("B.xlsx")}


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._handle_duplicate_columns")
@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._remove_unnamed_column_placeholders")
@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
def test_load(
    mock_excel_file: MagicMock,
    mock_remove_unnamed_column_placeholders: MagicMock,
    mock_handle_duplicate_columns: MagicMock,
    objects_excel: PandasExcelData,
):
    fs = ExcelFileStore(file_path=Path("input_data.xlsx"))
    mock_excel_file.return_value = MockExcelFile(objects_excel)
    mock_remove_unnamed_column_placeholders.side_effect = noop
    mock_handle_duplicate_columns.side_effect = noop

    # Act
    data = fs.load()

    # Assert
    mock_excel_file.assert_called_once()
    pd.testing.assert_frame_equal(data["Nodes"], objects_excel["Nodes"])
    pd.testing.assert_frame_equal(data["Lines"], objects_excel["Lines"])
    assert mock_remove_unnamed_column_placeholders.call_args_list[0] == call(data=objects_excel["Nodes"])
    assert mock_remove_unnamed_column_placeholders.call_args_list[1] == call(data=objects_excel["Lines"])
    assert mock_handle_duplicate_columns.call_args_list[0] == call(data=objects_excel["Nodes"], sheet_name="Nodes")
    assert mock_handle_duplicate_columns.call_args_list[1] == call(data=objects_excel["Lines"], sheet_name="Lines")


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._handle_duplicate_columns")
@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._remove_unnamed_column_placeholders")
@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
def test_load__extra(
    mock_excel_file: MagicMock,
    mock_remove_unnamed_column_placeholders: MagicMock,
    mock_handle_duplicate_columns: MagicMock,
    objects_excel: PandasExcelData,
    specs_excel: PandasExcelData,
):
    # Arrange
    fs = ExcelFileStore(Path("input_data.xlsx"), foo=Path("foo_types.xlsx"))
    mock_excel_file.side_effect = (MockExcelFile(objects_excel), MockExcelFile(specs_excel))
    mock_remove_unnamed_column_placeholders.side_effect = noop
    mock_handle_duplicate_columns.side_effect = noop

    # Act
    data = fs.load()

    # Assert
    assert mock_excel_file.call_count == 2
    pd.testing.assert_frame_equal(data["Nodes"], objects_excel["Nodes"])
    pd.testing.assert_frame_equal(data["Lines"], objects_excel["Lines"])
    pd.testing.assert_frame_equal(data["foo.Colors"], specs_excel["Colors"])
    pd.testing.assert_frame_equal(data["foo.Lines"], specs_excel["Lines"])
    assert mock_remove_unnamed_column_placeholders.call_args_list[0] == call(data=objects_excel["Nodes"])
    assert mock_remove_unnamed_column_placeholders.call_args_list[1] == call(data=objects_excel["Lines"])
    assert mock_remove_unnamed_column_placeholders.call_args_list[2] == call(data=specs_excel["Colors"])
    assert mock_remove_unnamed_column_placeholders.call_args_list[3] == call(data=specs_excel["Lines"])
    assert mock_handle_duplicate_columns.call_args_list[0] == call(data=objects_excel["Nodes"], sheet_name="Nodes")
    assert mock_handle_duplicate_columns.call_args_list[1] == call(data=objects_excel["Lines"], sheet_name="Lines")
    assert mock_handle_duplicate_columns.call_args_list[2] == call(data=specs_excel["Colors"], sheet_name="Colors")
    assert mock_handle_duplicate_columns.call_args_list[3] == call(data=specs_excel["Lines"], sheet_name="Lines")


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._handle_duplicate_columns")
@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._remove_unnamed_column_placeholders")
@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelFile")
def test_load__extra__duplicate_sheet_name(
    mock_excel_file: MagicMock,
    mock_remove_unnamed_column_placeholders: MagicMock,
    mock_handle_duplicate_columns: MagicMock,
):
    # Arrange
    foo_data = {"bar.Nodes": pd.DataFrame()}
    bar_data = {"Nodes": pd.DataFrame()}
    fs = ExcelFileStore(Path("foo.xlsx"), bar=Path("bar.xlsx"))
    mock_excel_file.side_effect = (MockExcelFile(foo_data), MockExcelFile(bar_data))
    mock_remove_unnamed_column_placeholders.side_effect = noop
    mock_handle_duplicate_columns.side_effect = noop

    # Act / Assert
    with pytest.raises(ValueError, match=r"Duplicate sheet name.+bar\.Nodes"):
        fs.load()


@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelWriter")
@patch("power_grid_model_io.data_stores.excel_file_store.pd.DataFrame.to_excel")
def test_save(mock_to_excel: MagicMock, mock_excel_writer: MagicMock):
    # Arrange
    data = TabularData(
        foo=pd.DataFrame([(1, 2.3)], columns=["id", "foo_val"]),
        bar=np.array([(4, 5.6)], dtype=[("id", "i4"), ("bar_val", "f4")]),
    )
    fs = ExcelFileStore(Path("output_data.xlsx"))

    # Act
    fs.save(data=data)

    # Assert
    mock_excel_writer.assert_called_once_with(path=Path("output_data.xlsx"))
    mock_to_excel.assert_any_call(excel_writer=mock_excel_writer().__enter__(), sheet_name="foo")
    mock_to_excel.assert_any_call(excel_writer=mock_excel_writer().__enter__(), sheet_name="bar")


@patch("power_grid_model_io.data_stores.excel_file_store.pd.ExcelWriter")
@patch("power_grid_model_io.data_stores.excel_file_store.pd.DataFrame.to_excel")
def test_save__multiple_files(mock_to_excel: MagicMock, mock_excel_writer):
    # Arrange
    data = TabularData(
        **{
            "nodes": pd.DataFrame(),
            "lines": pd.DataFrame(),
            "foo.colors": pd.DataFrame(),
        }
    )
    fs = ExcelFileStore(Path("output_data.xlsx"), foo=Path("foo.xlsx"), bar=Path("bar.xlsx"))

    output_data_writer = MagicMock()
    foo_writer = MagicMock()
    mock_excel_writer.side_effect = [output_data_writer, foo_writer]

    # Act
    fs.save(data=data)

    # Assert
    mock_excel_writer.assert_any_call(path=Path("output_data.xlsx"))
    mock_to_excel.assert_any_call(excel_writer=output_data_writer.__enter__(), sheet_name="nodes")
    mock_to_excel.assert_any_call(excel_writer=output_data_writer.__enter__(), sheet_name="lines")

    mock_excel_writer.assert_any_call(path=Path("foo.xlsx"))
    mock_to_excel.assert_any_call(excel_writer=foo_writer.__enter__(), sheet_name="colors")


@pytest.mark.parametrize(
    ("column_name", "is_unnamed"),
    [
        param("", False, id="empty"),
        param("id", False, id="id"),
        param("Unnamed: 123_level_1", True, id="unnamed"),
        param("Unnamed", False, id="not_unnamed"),
        param("123", False, id="number_as_str"),
    ],
)
def test_unnamed_pattern(column_name: str, is_unnamed: bool):
    assert bool(ExcelFileStore._unnamed_pattern.fullmatch(column_name)) == is_unnamed


def test_remove_unnamed_column_placeholders():
    # Arrange
    data = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["ID", "Unnamed: 123_level_0", "X"])
    store = ExcelFileStore()

    # Act
    result = store._remove_unnamed_column_placeholders(data=data)

    # Assert
    pd.testing.assert_frame_equal(result, pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["ID", "", "X"]))


def test_remove_unnamed_column_placeholders__multi_first():
    # Arrange
    columns = pd.MultiIndex.from_tuples([("ID", "A"), ("Unnamed: 123_level_0", "B"), ("X", "C")])
    data = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=columns)
    store = ExcelFileStore()

    # Act
    result = store._remove_unnamed_column_placeholders(data=data)

    # Assert
    columns = pd.MultiIndex.from_tuples([("ID", "A"), ("", "B"), ("X", "C")])
    pd.testing.assert_frame_equal(result, pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=columns))


def test_remove_unnamed_column_placeholders__multi_second():
    # Arrange
    columns = pd.MultiIndex.from_tuples([("ID", ""), ("B", "Unnamed: 123_level_1"), ("C", "kW")])
    data = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=columns)
    store = ExcelFileStore()

    # Act
    result = store._remove_unnamed_column_placeholders(data=data)

    # Assert
    columns = pd.MultiIndex.from_tuples([("ID", ""), ("B", ""), ("C", "kW")])
    pd.testing.assert_frame_equal(result, pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=columns))


def test_remove_unnamed_column_placeholders__empty():
    # Arrange
    data = pd.DataFrame()
    store = ExcelFileStore()

    # Act
    result = store._remove_unnamed_column_placeholders(data=data)

    # Assert
    pd.testing.assert_frame_equal(result, data)


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._check_duplicate_values")
def test_handle_duplicate_columns(mock_check_duplicate_values: MagicMock):
    # Arrange
    data = pd.DataFrame(
        [  # A    B    C    A    B    A
            # 0    1    2    3    4    5
            [101, 201, 301, 101, 201, 101],
            [102, 202, 302, 111, 202, 102],
            [103, 203, 303, 103, 203, 103],
        ],
        columns=["A", "B", "C", "A", "B", "A"],
    )
    store = ExcelFileStore()
    mock_check_duplicate_values.return_value = {3: "A_2", 4: "B_2", 5: "A_3"}

    # Act
    with capture_logs() as cap_log:
        actual = store._handle_duplicate_columns(data=data, sheet_name="foo")

    # Assert
    assert len(cap_log) == 3
    assert_log_exists(cap_log, "warning", "Column is renamed", col_name="A", new_name="A_2", col_idx=3)
    assert_log_exists(cap_log, "warning", "Column is renamed", col_name="B", new_name="B_2", col_idx=4)
    assert_log_exists(cap_log, "warning", "Column is renamed", col_name="A", new_name="A_3", col_idx=5)

    expected = pd.DataFrame(
        [  # A    B    C   A_2  B_2  A_3
            [101, 201, 301, 101, 201, 101],
            [102, 202, 302, 111, 202, 102],
            [103, 203, 303, 103, 203, 103],
        ],
        columns=["A", "B", "C", "A_2", "B_2", "A_3"],
    )
    pd.testing.assert_frame_equal(actual, expected)


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._check_duplicate_values")
def test_handle_duplicate_columns__multi(mock_check_duplicate_values: MagicMock):
    # Arrange
    data = pd.DataFrame(
        [  # A,1  B,2  C,3  A,1  B,2  A,1
            # 0    1    2    3    4    5
            [101, 201, 301, 101, 201, 101],
            [102, 202, 302, 111, 202, 102],
            [103, 203, 303, 103, 203, 103],
        ],
        columns=pd.MultiIndex.from_tuples([("A", 1), ("B", 2), ("C", 3), ("A", 1), ("B", 2), ("A", 1)]),
    )
    store = ExcelFileStore()
    mock_check_duplicate_values.return_value = {3: ("A_2", 1), 4: ("B_2", 2), 5: ("A_3", 1)}

    # Act
    with capture_logs() as cap_log:
        actual = store._handle_duplicate_columns(data=data, sheet_name="foo")

    # Assert
    assert len(cap_log) == 3
    assert_log_exists(cap_log, "warning", "Column is renamed", col_name=("A", 1), new_name=("A_2", 1), col_idx=3)
    assert_log_exists(cap_log, "warning", "Column is renamed", col_name=("B", 2), new_name=("B_2", 2), col_idx=4)
    assert_log_exists(cap_log, "warning", "Column is renamed", col_name=("A", 1), new_name=("A_3", 1), col_idx=5)

    expected = pd.DataFrame(
        [  # A,1  B,2  C,3 A_2,1 B_2,2 A_3,1
            [101, 201, 301, 101, 201, 101],
            [102, 202, 302, 111, 202, 102],
            [103, 203, 303, 103, 203, 103],
        ],
        columns=pd.MultiIndex.from_tuples([("A", 1), ("B", 2), ("C", 3), ("A_2", 1), ("B_2", 2), ("A_3", 1)]),
    )
    pd.testing.assert_frame_equal(actual, expected)


def test_handle_duplicate_columns__empty():
    # Arrange
    data = pd.DataFrame()
    store = ExcelFileStore()

    # Act
    result = store._handle_duplicate_columns(data=data, sheet_name="foo")

    # Assert
    pd.testing.assert_frame_equal(result, data)


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._group_columns_by_index")
def test_check_duplicate_values(mock_group_columns: MagicMock):
    # Arrange
    data = pd.DataFrame(
        [  # A,1  B,2  C,3  A,1  B,2  A,1
            # 0    1    2    3    4    5
            [101, 201, 301, 101, 201, 101],
            [102, 202, 302, 111, 202, 102],
            [103, 203, 303, 103, 203, 103],
        ],
        columns=["A", "B", "C", "A", "B", "A"],
    )
    store = ExcelFileStore()
    mock_group_columns.return_value = {"A": {0, 3, 5}, "B": {1, 4}, "C": {2}}

    # Act
    with capture_logs() as cap_log:
        result = store._check_duplicate_values(data=data, sheet_name="foo")

    # Assert
    assert len(cap_log) == 2
    assert_log_exists(
        cap_log, "error", "Found duplicate column names, with different data", col_name="A", col_idx=[0, 3, 5]
    )
    assert_log_exists(cap_log, "warning", "Found duplicate column names, with same data", col_name="B", col_idx=[1, 4])

    assert result == {3: "A_2", 4: "B_2", 5: "A_3"}


@patch("power_grid_model_io.data_stores.excel_file_store.ExcelFileStore._group_columns_by_index")
def test_check_duplicate_values__multi(mock_group_columns: MagicMock):
    # Arrange
    data = pd.DataFrame(
        [  # A,1  B,2  C,3  A,1  B,2  A,1
            # 0    1    2    3    4    5
            [101, 201, 301, 101, 201, 101],
            [102, 202, 302, 111, 202, 102],
            [103, 203, 303, 103, 203, 103],
        ],
        columns=pd.MultiIndex.from_tuples([("A", 1), ("B", 2), ("C", 3), ("A", 1), ("B", 2), ("A", 1)]),
    )
    store = ExcelFileStore()
    mock_group_columns.return_value = {("A", 1): {0, 3, 5}, ("B", 2): {1, 4}, ("C", 3): {2}}

    # Act
    with capture_logs() as cap_log:
        result = store._check_duplicate_values(data=data, sheet_name="foo")

    # Assert
    assert len(cap_log) == 2
    assert_log_exists(
        cap_log, "error", "Found duplicate column names, with different data", col_name=("A", 1), col_idx=[0, 3, 5]
    )
    assert_log_exists(
        cap_log, "warning", "Found duplicate column names, with same data", col_name=("B", 2), col_idx=[1, 4]
    )

    assert result == {3: ("A_2", 1), 4: ("B_2", 2), 5: ("A_3", 1)}


def test_group_columns_by_index():
    # Arrange
    data = pd.DataFrame(columns=["A", "B", "C", "A", "B", "A"])

    # Act
    grouped = ExcelFileStore._group_columns_by_index(data=data)

    # Assert
    assert grouped == {"A": {0, 3, 5}, "B": {1, 4}, "C": {2}}


def test_group_columns_by_index__multi():
    # Arrange
    data = pd.DataFrame(columns=pd.MultiIndex.from_tuples([("A", 1), ("B", 2), ("C", 3), ("A", 1), ("B", 2), ("A", 1)]))

    # Act
    grouped = ExcelFileStore._group_columns_by_index(data=data)

    # Assert
    assert grouped == {("A", 1): {0, 3, 5}, ("B", 2): {1, 4}, ("C", 3): {2}}
