import json
import os
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Any, NamedTuple
from unittest import mock

import numpy as np
import pandas as pd
import prophet
import pytest
import yaml
from packaging.version import Version
from prophet import Prophet

import mlflow
import mlflow.prophet
import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
from mlflow import pyfunc
from mlflow.models import Model, infer_signature
from mlflow.models.utils import _read_example, load_serving_example
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.utils.environment import _mlflow_conda_env
from mlflow.utils.model_utils import _get_flavor_configuration

from tests.helper_functions import (
    _assert_pip_requirements,
    _compare_conda_env_requirements,
    _compare_logged_code_paths,
    _is_available_on_pypi,
    _mlflow_major_version_string,
    assert_register_model_called_with_local_model_path,
    pyfunc_serve_and_score_model,
)

EXTRA_PYFUNC_SERVING_TEST_ARGS = (
    [] if _is_available_on_pypi("prophet") else ["--env-manager", "local"]
)


class DataGeneration:
    def __init__(self, **kwargs):
        self.shift = kwargs["shift"]
        self.start = datetime.strptime(kwargs["start"], "%Y-%M-%d")
        self.size = kwargs["size"]
        self.date_field = kwargs["date_field"]
        self.target_field = kwargs["target_field"]
        self.seasonal_period = kwargs["seasonal_period"]
        self.seasonal_freq = kwargs["seasonal_freq"]
        np.random.seed(42)

    def _period_gen(self):
        period = np.sin(np.arange(0, self.seasonal_period, self.seasonal_freq)) * 50 + 50
        return np.tile(
            period, int(np.ceil(self.size / (self.seasonal_period / self.seasonal_freq)))
        )[: self.size]

    def _generate_raw(self):
        base = np.random.lognormal(mean=2.0, sigma=0.92, size=self.size)
        seasonal = [
            np.polyval([-5.0, -1.0], x) for x in np.linspace(start=0, stop=2, num=self.size)
        ]
        return (
            np.linspace(start=45.0, stop=90.0, num=self.size) + base + seasonal + self._period_gen()
        )

    def _generate_linear_data(self):
        class DataStruct(NamedTuple):
            dates: Any
            series: Any

        series = self._generate_raw()
        date_ranges = np.arange(
            self.start, self.start + timedelta(days=self.size), timedelta(days=1)
        ).astype(date)
        return DataStruct(date_ranges, series)

    def _generate_shift_data(self):
        class DataStruct(NamedTuple):
            dates: Any
            series: Any

        raw = self._generate_raw()[: int(self.size * 0.6)]
        temperature = np.concatenate((raw, raw / 2.0)).ravel()[: self.size]
        date_ranges = np.arange(
            self.start, self.start + timedelta(days=self.size), timedelta(days=1)
        ).astype(date)
        return DataStruct(date_ranges, temperature)

    def _gen_series(self):
        if self.shift:
            return self._generate_shift_data()
        else:
            return self._generate_linear_data()

    def create_series_df(self):
        gen_data = self._gen_series()
        temporal_df = pd.DataFrame.from_records(gen_data).T
        temporal_df.columns = [self.date_field, self.target_field]
        return temporal_df


TEST_CONFIG = {
    "shift": False,
    "start": "2011-07-25",
    "size": 365 * 4,
    "seasonal_period": 7,
    "seasonal_freq": 0.1,
    "date_field": "ds",
    "target_field": "y",
}
FORECAST_HORIZON = 60
SEED = 98765
HORIZON_FIELD_NAME = "horizon"
TARGET_FIELD_NAME = "yhat"
DS_FORMAT = "%Y-%m-%dT%H:%M:%S"
INFER_FORMAT = "%Y-%m-%d %H:%M:%S"


class ModelWithSource(NamedTuple):
    model: Any
    data: Any


@pytest.fixture(scope="module")
def prophet_model():
    np.random.seed(SEED)
    data = DataGeneration(**TEST_CONFIG).create_series_df()
    model = Prophet().fit(data)
    return ModelWithSource(model, data)


@pytest.fixture
def model_path(tmp_path):
    return tmp_path.joinpath("model")


@pytest.fixture
def prophet_custom_env(tmp_path):
    conda_env = tmp_path.joinpath("conda_env.yml")
    _mlflow_conda_env(conda_env, additional_pip_deps=["prophet"])
    return conda_env


def future_horizon_df(model, horizon):
    return model.make_future_dataframe(periods=horizon)


def generate_forecast(model, horizon):
    return model.predict(model.make_future_dataframe(periods=horizon))[TARGET_FIELD_NAME]


def test_model_native_save_load(prophet_model, model_path):
    model = prophet_model.model
    mlflow.prophet.save_model(pr_model=model, path=model_path)
    loaded_model = mlflow.prophet.load_model(model_uri=model_path)

    np.testing.assert_array_equal(
        generate_forecast(model, FORECAST_HORIZON),
        loaded_model.predict(future_horizon_df(loaded_model, FORECAST_HORIZON))[TARGET_FIELD_NAME],
    )


def test_model_pyfunc_save_load(prophet_model, model_path):
    model = prophet_model.model
    mlflow.prophet.save_model(pr_model=model, path=model_path)
    loaded_pyfunc = pyfunc.load_model(model_uri=model_path)

    horizon_df = future_horizon_df(model, FORECAST_HORIZON)

    np.testing.assert_array_equal(
        generate_forecast(model, FORECAST_HORIZON),
        loaded_pyfunc.predict(horizon_df)[TARGET_FIELD_NAME],
    )


@pytest.mark.parametrize("use_signature", [True, False])
@pytest.mark.parametrize("use_example", [True, False])
def test_signature_and_examples_saved_correctly(
    prophet_model, model_path, use_signature, use_example
):
    data = prophet_model.data
    model = prophet_model.model
    horizon_df = future_horizon_df(model, FORECAST_HORIZON)
    signature_ = infer_signature(data, model.predict(horizon_df))
    signature = signature_ if use_signature else None
    if use_example:
        example = data[0:5].copy(deep=False)
        example["y"] = pd.to_numeric(example["y"])  # cast to appropriate precision
    else:
        example = None
    mlflow.prophet.save_model(model, path=model_path, signature=signature, input_example=example)
    mlflow_model = Model.load(model_path)
    if signature is None and example is None:
        assert mlflow_model.signature is None
    else:
        assert mlflow_model.signature == signature_
    if example is None:
        assert mlflow_model.saved_input_example_info is None
    else:
        r_example = _read_example(mlflow_model, model_path).copy(deep=False)
        r_example["ds"] = pd.to_datetime(r_example["ds"], format=DS_FORMAT)
        np.testing.assert_array_equal(r_example, example)


def test_model_load_from_remote_uri_succeeds(prophet_model, model_path, mock_s3_bucket):
    mlflow.prophet.save_model(pr_model=prophet_model.model, path=model_path)

    artifact_root = f"s3://{mock_s3_bucket}"
    artifact_path = "model"
    artifact_repo = S3ArtifactRepository(artifact_root)
    artifact_repo.log_artifacts(model_path, artifact_path=artifact_path)

    # NB: cloudpathlib would need to be used here to handle object store uri
    model_uri = os.path.join(artifact_root, artifact_path)
    reloaded_prophet_model = mlflow.prophet.load_model(model_uri=model_uri)
    np.testing.assert_array_equal(
        generate_forecast(prophet_model.model, FORECAST_HORIZON),
        generate_forecast(reloaded_prophet_model, FORECAST_HORIZON),
    )


@pytest.mark.parametrize("should_start_run", [True, False])
def test_prophet_log_model(prophet_model, tmp_path, should_start_run):
    try:
        if should_start_run:
            mlflow.start_run()
        artifact_path = "prophet"
        conda_env = tmp_path.joinpath("conda_env.yaml")
        _mlflow_conda_env(conda_env, additional_pip_deps=["pystan", "prophet"])

        model_info = mlflow.prophet.log_model(
            prophet_model.model, name=artifact_path, conda_env=str(conda_env)
        )
        reloaded_prophet_model = mlflow.prophet.load_model(model_uri=model_info.model_uri)

        np.testing.assert_array_equal(
            generate_forecast(prophet_model.model, FORECAST_HORIZON),
            generate_forecast(reloaded_prophet_model, FORECAST_HORIZON),
        )

        model_path = Path(_download_artifact_from_uri(artifact_uri=model_info.model_uri))
        model_config = Model.load(str(model_path.joinpath("MLmodel")))
        assert pyfunc.FLAVOR_NAME in model_config.flavors
        assert pyfunc.ENV in model_config.flavors[pyfunc.FLAVOR_NAME]
        env_path = model_config.flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV]["conda"]
        assert model_path.joinpath(env_path).exists()

    finally:
        mlflow.end_run()


def test_log_model_calls_register_model(prophet_model, tmp_path):
    artifact_path = "prophet"
    register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model")
    with mlflow.start_run(), register_model_patch:
        conda_env = tmp_path.joinpath("conda_env.yaml")
        _mlflow_conda_env(conda_env, additional_pip_deps=["pystan", "prophet"])
        model_info = mlflow.prophet.log_model(
            prophet_model.model,
            name=artifact_path,
            conda_env=str(conda_env),
            registered_model_name="ProphetModel1",
        )
        assert_register_model_called_with_local_model_path(
            register_model_mock=mlflow.tracking._model_registry.fluent._register_model,
            model_uri=model_info.model_uri,
            registered_model_name="ProphetModel1",
        )


def test_log_model_no_registered_model_name(prophet_model, tmp_path):
    artifact_path = "prophet"
    register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model")
    with mlflow.start_run(), register_model_patch:
        conda_env = tmp_path.joinpath("conda_env.yaml")
        _mlflow_conda_env(conda_env, additional_pip_deps=["pystan", "prophet"])
        mlflow.prophet.log_model(prophet_model.model, name=artifact_path, conda_env=str(conda_env))
        mlflow.tracking._model_registry.fluent._register_model.assert_not_called()


def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
    prophet_model, model_path, prophet_custom_env
):
    mlflow.prophet.save_model(
        pr_model=prophet_model.model, path=model_path, conda_env=str(prophet_custom_env)
    )
    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    saved_conda_env_path = model_path.joinpath(pyfunc_conf[pyfunc.ENV]["conda"])

    assert saved_conda_env_path.exists()
    assert not prophet_custom_env.samefile(saved_conda_env_path)

    prophet_custom_env_parsed = yaml.safe_load(prophet_custom_env.read_bytes())
    saved_conda_env_parsed = yaml.safe_load(saved_conda_env_path.read_bytes())
    assert prophet_custom_env_parsed == saved_conda_env_parsed


def test_model_save_persists_requirements_in_mlflow_model_directory(
    prophet_model, model_path, prophet_custom_env
):
    mlflow.prophet.save_model(
        pr_model=prophet_model.model, path=model_path, conda_env=str(prophet_custom_env)
    )

    saved_pip_req_path = model_path.joinpath("requirements.txt")
    _compare_conda_env_requirements(prophet_custom_env, str(saved_pip_req_path))


def test_log_model_with_pip_requirements(prophet_model, tmp_path):
    expected_mlflow_version = _mlflow_major_version_string()
    req_file = tmp_path.joinpath("requirements.txt")
    req_file.write_text("a")
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name="model", pip_requirements=str(req_file)
        )
        _assert_pip_requirements(model_info.model_uri, [expected_mlflow_version, "a"], strict=True)

    # List of requirements
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name="model", pip_requirements=[f"-r {req_file}", "b"]
        )
        _assert_pip_requirements(
            model_info.model_uri, [expected_mlflow_version, "a", "b"], strict=True
        )

    # Constraints file
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name="model", pip_requirements=[f"-c {req_file}", "b"]
        )
        _assert_pip_requirements(
            model_info.model_uri,
            [expected_mlflow_version, "b", "-c constraints.txt"],
            ["a"],
            strict=True,
        )


def test_log_model_with_extra_pip_requirements(prophet_model, tmp_path):
    expected_mlflow_version = _mlflow_major_version_string()
    default_reqs = mlflow.prophet.get_default_pip_requirements()

    # Path to a requirements file
    req_file = tmp_path.joinpath("requirements.txt")
    req_file.write_text("a")
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name="model", extra_pip_requirements=str(req_file)
        )
        _assert_pip_requirements(
            model_info.model_uri, [expected_mlflow_version, *default_reqs, "a"]
        )

    # List of requirements
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name="model", extra_pip_requirements=[f"-r {req_file}", "b"]
        )
        _assert_pip_requirements(
            model_info.model_uri, [expected_mlflow_version, *default_reqs, "a", "b"]
        )

    # Constraints file
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name="model", extra_pip_requirements=[f"-c {req_file}", "b"]
        )
        _assert_pip_requirements(
            model_uri=model_info.model_uri,
            requirements=[expected_mlflow_version, *default_reqs, "b", "-c constraints.txt"],
            constraints=["a"],
            strict=False,
        )


def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
    prophet_model, model_path
):
    mlflow.prophet.save_model(prophet_model.model, model_path)
    _assert_pip_requirements(model_path, mlflow.prophet.get_default_pip_requirements())


def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
    prophet_model,
):
    artifact_path = "model"
    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(prophet_model.model, name=artifact_path)
    _assert_pip_requirements(model_info.model_uri, mlflow.prophet.get_default_pip_requirements())


def test_pyfunc_serve_and_score(prophet_model):
    artifact_path = "model"
    # cast to string representation of datetime series, otherwise will default cast to Unix time
    # which Prophet does not support for encoding
    inference_data = (
        prophet_model.model
        .make_future_dataframe(FORECAST_HORIZON)["ds"]
        .dt.strftime(INFER_FORMAT)
        .to_frame(name="ds")
    )
    with mlflow.start_run():
        extra_pip_requirements = (
            ["holidays<=0.24"] if Version(prophet.__version__) <= Version("1.1.3") else []
        ) + (["pandas<2"] if Version(prophet.__version__) < Version("1.1") else [])
        model_info = mlflow.prophet.log_model(
            prophet_model.model,
            name=artifact_path,
            extra_pip_requirements=extra_pip_requirements,
            input_example=inference_data,
        )
    local_predict = prophet_model.model.predict(
        prophet_model.model.make_future_dataframe(FORECAST_HORIZON)
    )

    inference_payload = load_serving_example(model_info.model_uri)
    resp = pyfunc_serve_and_score_model(
        model_info.model_uri,
        data=inference_payload,
        content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
        extra_args=EXTRA_PYFUNC_SERVING_TEST_ARGS,
    )
    scores = pd.DataFrame(data=json.loads(resp.content.decode("utf-8"))["predictions"])

    # predictions are deterministic, but yhat_lower, yhat_upper are non-deterministic based on
    # stan build underlying environment. Seed value only works for reproducibility of yhat.
    # see: https://github.com/facebook/prophet/issues/1124
    pd.testing.assert_series_equal(
        left=local_predict["yhat"], right=scores["yhat"], check_dtype=True
    )


def test_log_model_with_code_paths(prophet_model):
    artifact_path = "model"
    with (
        mlflow.start_run(),
        mock.patch("mlflow.prophet._add_code_from_conf_to_system_path") as add_mock,
    ):
        model_info = mlflow.prophet.log_model(
            prophet_model.model, name=artifact_path, code_paths=[__file__]
        )
        _compare_logged_code_paths(__file__, model_info.model_uri, mlflow.prophet.FLAVOR_NAME)
        mlflow.prophet.load_model(model_info.model_uri)
        add_mock.assert_called()


def test_virtualenv_subfield_points_to_correct_path(prophet_model, model_path):
    mlflow.prophet.save_model(prophet_model.model, path=model_path)
    pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
    python_env_path = Path(model_path, pyfunc_conf[pyfunc.ENV]["virtualenv"])
    assert python_env_path.exists()
    assert python_env_path.is_file()


def test_model_save_load_with_metadata(prophet_model, model_path):
    mlflow.prophet.save_model(
        prophet_model.model, path=model_path, metadata={"metadata_key": "metadata_value"}
    )

    reloaded_model = mlflow.pyfunc.load_model(model_uri=model_path)
    assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value"


def test_model_log_with_metadata(prophet_model):
    artifact_path = "model"

    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(
            prophet_model.model,
            name=artifact_path,
            metadata={"metadata_key": "metadata_value"},
        )

    reloaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
    assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value"


def test_model_log_with_signature_inference(prophet_model):
    artifact_path = "model"
    model = prophet_model.model
    horizon_df = future_horizon_df(model, FORECAST_HORIZON)
    signature = infer_signature(horizon_df, model.predict(horizon_df))

    with mlflow.start_run():
        model_info = mlflow.prophet.log_model(model, name=artifact_path, input_example=horizon_df)

    loaded_model = Model.load(model_info.model_uri)
    assert loaded_model.signature == signature