import io
import json
import time
import warnings
from datetime import datetime, timedelta
from pprint import pprint

import pandas as pd
import requests
from typeguard import typechecked

from ._version import __version__
from .settings import ValidStatus

ALLOWED_DATAFRAME_SIZE = 5.5 * int(1e6)

PING_TIME_INITIAL = 1.0  # Seconds
PING_FRACTIONAL_INCREASE = 0.1

DATETIME_STRING_FORMAT = "%Y-%m-%d %H:%M:%S"


@typechecked
def _calculate_ping_time(elapsed_time: float) -> float:
    # This smoothly transitions between regular pinging at the initial ping time
    # to more drawn out pinging (expoential) as time goes on
    return PING_TIME_INITIAL + elapsed_time * PING_FRACTIONAL_INCREASE


@typechecked
def wait_for_job_completion(
    api_function: callable, *args, verbose: bool = False
) -> dict:
    start_time = time.time()
    status = ValidStatus.PROCESSING
    while status == ValidStatus.PROCESSING:
        elapsed_time = time.time() - start_time  # This will be ~0 seconds initially
        wait_time = _calculate_ping_time(elapsed_time)
        time.sleep(wait_time)
        _, response = api_function(*args)
        status = ValidStatus(get_value_from_body("status", response))
        if verbose:
            t = timedelta(seconds=elapsed_time)
            t -= timedelta(microseconds=t.microseconds)  # Remove microseconds
            print(f"{str(t)}: Job status: {status.value}")
    return response


@typechecked
def calculate_run_time(start_time_str: str, end_time_str: str) -> str:
    # convert the start_time and end_time to datetime
    start_time = datetime.fromisoformat(start_time_str)
    end_time = datetime.fromisoformat(end_time_str)

    # calculate the run time
    run_time = end_time - start_time

    # round the run time to the nearest second
    run_time_seconds = run_time.total_seconds()
    run_time_rounded = round(run_time_seconds)
    run_time = timedelta(seconds=run_time_rounded)

    # format the run time in a nice format
    run_time_nice_format = str(run_time)

    return run_time_nice_format


@typechecked
def get_message(response: dict) -> str:
    # TODO: This could be a method of the response object
    # TODO: This should be better
    try:
        message = response["message"]
    except:
        message = response
    return message


@typechecked
def get_value_from_body(key: str, body: dict):
    # Relat responses from api.py directly
    # This improves error messaging
    if key in body.keys():
        return body[f"{key}"]
    else:
        raise KeyError(f"{key} not in API response body")


@typechecked
def calculate_runtime(start_time: str, end_time: str) -> str:
    # Calculate the runtime of a job given the start and end times.

    start_datetime = datetime.fromisoformat(start_time)
    end_datetime = datetime.fromisoformat(end_time)
    runtime = end_datetime - start_datetime

    # Extract hours, minutes, and seconds from the timedelta object
    total_seconds = int(runtime.total_seconds())
    hours, remainder = divmod(total_seconds, 3600)
    minutes, seconds = divmod(remainder, 60)

    # Format the runtime as HH:MM:SS
    nice_time_string = f"{hours:02}:{minutes:02}:{seconds:02}"

    return nice_time_string


@typechecked
def convert_time_format(time_string: str) -> str:
    # Parse the time string into a datetime object
    dt = datetime.fromisoformat(time_string)

    # Format the datetime object into a nice string
    nice_time_string = dt.strftime(DATETIME_STRING_FORMAT)

    return nice_time_string


@typechecked
def convert_time_formats_in_status(status: dict) -> dict:

    start_time = status.get("start_time", None)
    if start_time:
        status["start_time"] = convert_time_format(start_time)
    end_time = status.get("end_time", None)
    if end_time:
        status["end_time"] = convert_time_format(end_time)
    return status


@typechecked
def check_dataset(string: str) -> None:
    # Check that a sensible dataframe can be created from a .csv file string.

    # Check for duplicate columns
    # TODO: This assumes label is row 0
    header = pd.read_csv(io.StringIO(string), header=None, nrows=1).iloc[0].to_list()
    if len(set(header)) != len(header):
        raise TypeError("Dataset must contain no duplicate column names.")

    string_io = io.StringIO(string)
    try:
        df = pd.read_csv(string_io)
    except Exception:
        raise TypeError("Could not parse the input into a dataframe.")

    # Check that dataset has at least one column.
    if df.shape[0] < 1:
        raise TypeError("Dataset must have at least one column.")

    # Check that dataset has no duplicate column names.
    # TODO: Is this needed? What if the columns with identical names are not used in training?
    if len(set(df.columns)) != len(df.columns):
        raise TypeError(
            "Dataset must contain no duplicate column names."
        )  # Unable to raise this error as column names when read from a string are not recognised?

    # Check that the dataset contains only numerical values.
    if not df.map(lambda x: isinstance(x, (int, float))).all().all():
        raise Warning("Dataset contains non-numerical values.")


@typechecked
def download_file_from_url(presigned_url: str, file_path: str) -> None:
    try:
        # Make the HTTP GET request to fetch the file
        response = requests.get(presigned_url)

        # Write the content to the specified file path
        with open(file_path, "wb") as file:
            file.write(response.content)

        print(f"File downloaded successfully and saved to {file_path}")

    except Exception as e:
        print(f"Failed to download the file: {e}")


@typechecked
def upload_dataframe_to_presigned_url(
    df: pd.DataFrame,
    url: str,
    verbose: bool = False,
    check: bool = False,
) -> None:

    if check:
        csv_string = df.to_csv(index=False)
        check_dataset(csv_string)
    headers = {"Content-Type": "application/octet-stream"}

    # Create a buffer
    buffer = io.StringIO()

    # Write DataFrame to the buffer
    df.to_csv(buffer, index=False)

    # Convert buffer to string
    csv_string = buffer.getvalue()
    upload_file = json.dumps({"dataset": csv_string})
    response = requests.put(url, data=upload_file, headers=headers)
    if verbose:
        if response.status_code == 200:
            print(f"Dataframe is uploading.")
        else:
            print(f"Dataframe upload failed")
            print(f"Status code: {response.status_code}")
            print(f"Reason: {response.text}")


@typechecked
def upload_file_to_presigned_url(
    file_path: str,
    url: str,
    verbose: bool = False,
    check: bool = False,
) -> None:
    # Upload a file to the specified pre-signed URL.
    if check:
        with open(file_path, "rb") as file:
            csv_string = file.read().decode("utf-8")
            check_dataset(csv_string)
    with open(file_path, "rb") as file:
        headers = {"Content-Type": "application/octet-stream"}
        response = requests.put(url, data=file, headers=headers)
    if verbose:
        if response.status_code == 200:
            print(f"File {file_path} is uploading.")
        else:
            print(f"File upload failed")
            print(f"Status code: {response.status_code}")
            print(f"Reason: {response.text}")


@typechecked
def process_dataset_response(response: dict) -> pd.DataFrame:
    # Process the response from the API into a `pandas.DataFrame`.
    if response.get("dataset") is not None:
        csv = response["dataset"]
        csv = io.StringIO(csv)
    elif response.get("dataset_url") is not None:
        url = response["dataset_url"]
        csv = download_dataframe_from_presigned_url(url)
    else:
        raise ValueError("No 'dataset' or 'dataset_url' in the response.")
    df = pd.read_csv(csv, sep=",")
    return df


def process_result_response(response: dict) -> dict:
    # Process the response from the API into a `pandas.DataFrame`.
    if response.get("result") is not None:
        result = response["result"]
    elif response.get("result_url") is not None:
        url = response["result_url"]
        result = download_result_from_presigned_url(url)
        result = json.loads(result)
    else:
        raise ValueError("No 'result' or 'result_url' in the response.")
    return result


@typechecked
def download_result_from_presigned_url(url: str) -> io.StringIO:
    # Download a `pandas.DataFrame` from the specified pre-signed URL.
    response = requests.get(url)
    result = response.json()
    return result


@typechecked
def download_dataframe_from_presigned_url(url: str) -> io.StringIO:
    # Download a `pandas.DataFrame` from the specified pre-signed URL.
    response = requests.get(url)
    buffer = io.StringIO(response.json()["dataset"])
    return buffer


@typechecked
def get_csv_string(df: pd.DataFrame) -> str:
    buffer = io.StringIO()
    df.to_csv(buffer, index=False)
    return buffer.getvalue()


@typechecked
def remove_none_values(d: dict) -> dict:
    # Remove None values from a nested dictionary.
    return {
        k: remove_none_values(v) if isinstance(v, dict) else v
        for k, v in d.items()
        if v is not None
    }


class EmulatorResultsAdapter:
    # Utility classes for munging the results of an emulator API call.

    def __init__(self, method: str, response: dict):
        self.method = method
        self.status = ValidStatus(
            response["status"]
        )  # Set the self.status as the status of the response converted to a ValidStatus
        self.result = response[
            "result"
        ]  # Set the self.result as the result of the response
        self.result_url = response[
            "result_url"
        ]  # Set the self.result_url as the result_url of the response

    def _check_status(self):
        if self.status == ValidStatus.SUCCESS:
            return True
        else:
            return False

    def _get_result(self):
        if self.result:
            # If the result is already in the response, do nothing
            pass
        elif self.result_url:
            # If a result_url is in the response, download the result
            response = requests.get(self.result_url)
            # Set the result to the downloaded result
            self.result = response.json()
        else:
            raise ValueError("No result or result_url in the response.")

    def adapt_result(self, *args, verbose: bool = False):
        # Only run if the status is success
        if self._check_status():

            self._get_result()

            # Run the adapter method for the specific method
            method = getattr(self, f"adapt_{self.method}_results")

            return method(*args, verbose=verbose)
        else:
            # What should this return?
            raise RuntimeError("Non-success status")

    def adapt_update_results(self, verbose: bool):
        return self.result["update"]

    def adapt_design_results(self, verbose: bool):
        design_csv = self.result["design"]

        design = pd.read_csv(io.StringIO(design_csv))
        return design

    def adapt_predict_results(self, verbose: bool):
        mean_csv = self.result["mean"]
        std_csv = self.result["std"]

        mean = pd.read_csv(io.StringIO(mean_csv))
        std = pd.read_csv(io.StringIO(std_csv))

        if verbose:
            print("Mean predictions:")
            pprint(mean)
            print("Standard deviation predictions:")
            pprint(std)
        return mean, std

    def adapt_sample_results(self, verbose: bool):
        sample_csv = self.result["samples"]
        sample = pd.read_csv(io.StringIO(sample_csv), header=[0, 1], sep=",")

        if verbose:
            print("Samples:")
            pprint(sample)
        return sample

    def adapt_recommend_results(self, verbose: bool):
        recommendation_csv = self.result["candidate_points"]
        acq_func_value = self.result["acq_func"]

        recommendation = pd.read_csv(io.StringIO(recommendation_csv))

        if verbose:
            print("Recommended points:")
            pprint(recommendation)
        return recommendation, acq_func_value

    def adapt_maximize_results(self, verbose: bool):
        recommendation_csv = self.result["max_point"]

        max_point = pd.read_csv(io.StringIO(recommendation_csv))

        if verbose:
            print("Maximimum point:")
            pprint(max_point)
        return max_point

    def adapt_calibrate_results(self, verbose: bool):
        calibrated_csv = self.result["calibration"]
        calibration = pd.read_csv(io.StringIO(calibrated_csv), sep=",")

        # TODO: This is a nasty hack for when return_summary=True
        if "Unnamed: 0" in calibration.columns:
            calibration.set_index("Unnamed: 0", inplace=True)
            calibration.index.name = None

        if verbose:
            print("Calibration:")
            pprint(calibration)

        return calibration

    def adapt_score_results(self, verbose: bool):
        score = self.result["score"]

        if score:
            if isinstance(score, str):
                score = pd.read_csv(io.StringIO(score))
            else:
                pass
        else:
            warnings.warn(
                "No test data was available for this emulator, so it cannot be scored."
            )

        if verbose:
            print("Emulator score:")
            pprint(score)
        return score

    def adapt_benchmark_results(self, verbose: bool):
        benchmark = self.result["calibration_curve"]

        if benchmark:
            benchmark = pd.read_csv(io.StringIO(benchmark))
        else:
            warnings.warn(
                "No test data was available for this emulator, so it cannot be scored."
            )

        if verbose:
            print("Calibration curve:")
            pprint(benchmark)
        return benchmark

    def adapt_export_results(self, file_path: str, verbose: bool):

        download_file_from_url(self.result_url, file_path)


@typechecked
def reformat_summary_dict(summary_dict: dict, detailed: bool = False) -> dict:
    # Function to reformat the summary dictionary.
    properties, mean, kernel, likelihood = {}, {}, {}, {}

    # To retrieve the dictionary from the "base_estimator_diagnostics" key, when input or output decomposition is applied
    if "base_estimator_diagnostics" in summary_dict["estimator_diagnostics"].keys():
        estimator_diagnostics = summary_dict["estimator_diagnostics"].get(
            "base_estimator_diagnostics"
        )
    else:
        estimator_diagnostics = summary_dict.get("estimator_diagnostics")

    # Non-detailed case when only the learned hyperparameters are returned along with the covar_module and mean_module
    if "learned_hyperparameters" in estimator_diagnostics.keys():
        learned_params = estimator_diagnostics.get("learned_hyperparameters")

        # Augment properties with parameters of the variational distribution for Variational GP
        if "variational_strategy.inducing_points" in learned_params.keys():
            properties["inducing_points"] = learned_params.get(
                "variational_strategy.inducing_points"
            )
            properties["variational_distribution_mean"] = learned_params.get(
                "variational_strategy._variational_distribution.variational_mean"
            )
            properties["variational_distribution_covar"] = learned_params.get(
                "variational_strategy._variational_distribution.chol_variational_covar"
            )

        # Write the standard learned hyperparameters to the properties dictionary
        properties["covariance_noise"] = learned_params.get(
            "likelihood.noise_covar.original_noise"
        )
        mean["mean"] = learned_params.get("mean_module.original_constant")
        mean["mean_function_used"] = estimator_diagnostics.get("mean_module")
        kernel["lengthscale"] = learned_params.get(
            "covar_module.base_kernel.original_lengthscale"
        )
        kernel["outputscale"] = learned_params.get("covar_module.original_outputscale")
        kernel["kernel_function_used"] = estimator_diagnostics["covar_module"].replace(
            "\n", ""
        )
    else:
        return summary_dict
    if not detailed:
        return {"properties": properties, "mean": mean, "kernel": kernel}
    else:
        transform_list = ["input_transform_parameters", "outcome_transform_parameters"]
        join_delimiter = "_"

        # Extracting transform parameters
        if (
            transform_list[0] in estimator_diagnostics.keys()
            and transform_list[1] in estimator_diagnostics.keys()
        ):
            for transform_parameters in transform_list:
                for key in estimator_diagnostics[transform_parameters].keys():
                    # To avoid double underscores for some transform parameters
                    param_name = key.split(".")
                    if param_name[1][0] == "_":
                        param_name = param_name[0] + param_name[1]
                    else:
                        param_name = param_name[0] + "_" + param_name[1]
                    properties[param_name] = estimator_diagnostics[
                        transform_parameters
                    ].get(key)
            properties.pop(
                "outcome_transform_is_trained", None
            )  # Remove outcome_transform_is_trained parameter

        # Extract mean parameters
        if "mean_module_parameters" in estimator_diagnostics.keys():
            mean["raw_constant"] = estimator_diagnostics["mean_module_parameters"].get(
                "mean_module.raw_constant"
            )

        # Extracting kernel parameters
        if "covar_module_parameters" in estimator_diagnostics.keys():
            covar_module_params = estimator_diagnostics["covar_module_parameters"]
            for key in covar_module_params.keys():
                param_name = key.split(".")
                # Retrieving parameter names for Multi-Fidelity GP
                if len(param_name) > 5 and param_name[3] == "0":
                    param_name = join_delimiter.join(param_name[4:])
                # Retrieving parameter names for Mixed Single Task GP
                elif len(param_name) > 6 and param_name[1] == "kernels":
                    param_name = join_delimiter.join(param_name[2:])
                else:
                    del param_name[0]
                    param_name = join_delimiter.join(param_name)
                kernel[param_name] = covar_module_params.get(key)

        # Extracting likelihood parameters
        if "likelihood_parameters" in estimator_diagnostics.keys():
            likelihood_params = estimator_diagnostics["likelihood_parameters"]
            for key in likelihood_params.keys():
                param_name = key.split(".")
                # Shorten names differently for heteroskedastic noise model and other models
                if len(param_name) < 5:
                    param_name = join_delimiter.join(param_name[2:])
                else:
                    param_name = join_delimiter.join(param_name[3:])
                if (
                    param_name[0] == "_"
                ):  # To avoid double underscores for some likelihood parameters
                    likelihood["noise_model" + param_name] = likelihood_params.get(key)
                else:
                    likelihood["noise_model_" + param_name] = likelihood_params.get(key)
        return {
            "properties": properties,
            "mean": mean,
            "kernel": kernel,
            "likelihood": likelihood,
        }
