import csv
import datetime
import hashlib
import numpy as np
import os
import pandas as pd
import json

import verboselogs, logging

logger = verboselogs.VerboseLogger(__name__)

import random
from itertools import chain
from functools import partial

import postgres_copy
import sqlalchemy
from retrying import retry
from sqlalchemy.orm import sessionmaker
from ohio import PipeTextIO

from triage.component.results_schema import (
    Experiment,
    Matrix,
    Model,
    ExperimentMatrix,
    ExperimentModel,
    TriageRun,
)


def filename_friendly_hash(inputs):
    def dt_handler(x):
        if isinstance(x, datetime.datetime) or isinstance(x, datetime.date):
            return x.isoformat()
        raise TypeError("Unknown type")

    return hashlib.md5(
        json.dumps(inputs, default=dt_handler, sort_keys=True).encode("utf-8")
    ).hexdigest()


def get_subset_table_name(subset_config):
    return "subset_{}_{}".format(
        subset_config.get("name", "default"),
        filename_friendly_hash(subset_config),
    )


def retry_if_db_error(exception):
    return isinstance(exception, sqlalchemy.exc.OperationalError)


DEFAULT_RETRY_KWARGS = {
    "retry_on_exception": retry_if_db_error,
    "wait_exponential_multiplier": 1000,  # wait 2^x*1000ms between each retry
    "stop_max_attempt_number": 14,
    # with this configuration, last wait will be ~2 hours
    # for a total of ~4.5 hours waiting
}


db_retry = retry(**DEFAULT_RETRY_KWARGS)


@db_retry
def save_experiment_and_get_hash(config, db_engine):
    experiment_hash = filename_friendly_hash(config)
    session = sessionmaker(bind=db_engine)()
    session.merge(Experiment(experiment_hash=experiment_hash, config=config))
    session.commit()
    session.close()
    return experiment_hash


@db_retry
def associate_matrices_with_experiment(experiment_hash, matrix_uuids, db_engine):
    session = sessionmaker(bind=db_engine)()
    for matrix_uuid in matrix_uuids:
        session.merge(
            ExperimentMatrix(experiment_hash=experiment_hash, matrix_uuid=matrix_uuid)
        )
    session.commit()
    session.close()
    logger.spam("Associated matrices with experiment in database")


@db_retry
def associate_models_with_experiment(experiment_hash, model_hashes, db_engine):
    session = sessionmaker(bind=db_engine)()
    for model_hash in model_hashes:
        session.merge(
            ExperimentModel(experiment_hash=experiment_hash, model_hash=model_hash)
        )
    session.commit()
    session.close()
    logger.spam("Associated models with experiment in database")


@db_retry
def missing_matrix_uuids(experiment_hash, db_engine):
    """Compare the contents of the experiment_matrices table with that of the
    matrices table to produce a list of matrix_uuids that the experiment wants
    but are not available.
    """
    query = f"""
        select experiment_matrices.matrix_uuid
        from {ExperimentMatrix.__table__.fullname} experiment_matrices
        left join {Matrix.__table__.fullname} matrices
        on (experiment_matrices.matrix_uuid = matrices.matrix_uuid)
        where experiment_hash = %s
        and matrices.matrix_uuid is null
    """
    return [row[0] for row in db_engine.execute(query, experiment_hash)]


@db_retry
def missing_model_hashes(experiment_hash, db_engine):
    """Compare the contents of the experiment_models table with that of the
    models table to produce a list of model hashes the experiment wants
    but are not available.
    """
    query = f"""
        select experiment_models.model_hash
        from {ExperimentModel.__table__.fullname} experiment_models
        left join {Model.__table__.fullname} models
        on (experiment_models.model_hash = models.model_hash)
        where experiment_hash = %s
        and models.model_hash is null
    """
    return [row[0] for row in db_engine.execute(query, experiment_hash)]


class Batch:
    # modified from
    # http://codereview.stackexchange.com/questions/118883/split-up-an-iterable-into-batches
    def __init__(self, iterable, limit=None):
        self.iterator = iter(iterable)
        self.limit = limit
        try:
            self.current = next(self.iterator)
        except StopIteration:
            self.on_going = False
        else:
            self.on_going = True

    def group(self):
        yield self.current
        # start enumerate at 1 because we already yielded the last saved item
        for num, item in enumerate(self.iterator, 1):
            self.current = item
            if num == self.limit:
                break
            yield item
        else:
            self.on_going = False

    def __iter__(self):
        while self.on_going:
            yield self.group()


AVAILABLE_TIEBREAKERS = {"random", "best", "worst"}


def sort_predictions_and_labels(
    predictions_proba, labels, df_index, tiebreaker="random", sort_seed=None
):
    """Sort predictions and labels with a configured tiebreaking rule

    Args:
        predictions_proba (np.array) The predicted scores
        labels (np.array) The numeric labels (1/0, not True/False)
        df_index (pd.MultiIndex) Index (generally entity_id, as_of_date tuples) to be sorted with the labels/scores
        tiebreaker (string) The tiebreaking method ('best', 'worst', 'random')
        sort_seed (signed int) The sort seed. Needed if 'random' tiebreaking is picked.

    Returns:
        (tuple) (predictions_proba, labels, df_index), sorted
    """
    if len(labels) == 0:
        logger.notice("No labels present, skipping predictions sorting .")
        return (predictions_proba, labels, df_index)

    df = pd.DataFrame(predictions_proba, columns=["score"])
    df["label_value"] = labels
    df.set_index(df_index, inplace=True)

    if tiebreaker == "random":
        if not sort_seed:
            raise ValueError("If random tiebreaker is used, a sort seed must be given")
        random.seed(sort_seed)
        np.random.seed(sort_seed)
        df["random"] = np.random.rand(len(df))
        df.sort_values(by=["score", "random"], inplace=True, ascending=[False, False])
        df.drop("random", axis=1)
    elif tiebreaker == "worst":
        df.sort_values(
            by=["score", "label_value"],
            inplace=True,
            ascending=[False, True],
            na_position="first",
        )
    elif tiebreaker == "best":
        df.sort_values(
            by=["score", "label_value"],
            inplace=True,
            ascending=[False, False],
            na_position="last",
        )
    else:
        raise ValueError(f"Unknown tiebreaker: {tiebreaker}")

    return [df["score"].to_numpy(), df["label_value"].to_numpy(), df.index]


@db_retry
def retrieve_model_id_from_hash(db_engine, model_hash):
    """Retrieves a model id from the database that matches the given hash

    Args:
        db_engine (sqlalchemy.engine) A database engine
        model_hash (str) The model hash to lookup

    Returns: (int) The model id (if found in DB), None (if not)
    """
    session = sessionmaker(bind=db_engine)()
    try:
        saved = session.query(Model).filter_by(model_hash=model_hash).one_or_none()
        return saved.model_id if saved else None
    finally:
        session.close()


@db_retry
def retrieve_model_hash_from_id(db_engine, model_id):
    """Retrieves the model hash associated with a given model id

    Args:
        model_id (int) The id of a given model in the database

    Returns: (str) the stored hash of the model
    """
    session = sessionmaker(bind=db_engine)()
    try:
        return session.query(Model).get(model_id).model_hash
    finally:
        session.close()


@db_retry
def retrieve_existing_model_random_seeds(
    db_engine,
    model_group_id,
    train_end_time,
    train_matrix_uuid,
    training_label_timespan,
    experiment_random_seed,
):
    """Retrieve existing model random seeds matching the model parameters and
    experiment-level random seed to allow for reusing seeds before creating a
    new one.
    """
    query = f"""
        select models.random_seed
        from {ExperimentModel.__table__.fullname} experiment_models
        join {Model.__table__.fullname} models
        on (experiment_models.model_hash = models.model_hash)
        join {TriageRun.__table__.fullname} triage_runs
        on (
            experiment_models.experiment_hash = triage_runs.run_hash
            --makes sure that the model has been created before, and not being created in this run
            and models.built_in_triage_run = triage_runs.id
        )
        where models.model_group_id = {model_group_id}
        and models.train_end_time = '{train_end_time}'
        and models.train_matrix_uuid = '{train_matrix_uuid}'
        and models.training_label_timespan = '{training_label_timespan}'
        and triage_runs.random_seed = {experiment_random_seed}
        order by models.run_time DESC, random()
    """
    logging.debug(f"Query that will retrieve random seeds for the model: {query}")
    return [
        row[0]
        for row in db_engine.execute(query)
    ]


@db_retry
def retrieve_experiment_seed_from_run_id(db_engine, run_id):
    """Retrieves the random seed associated with a given experiment run

    Args:
        run_id (int) The id of a given experiment run in the database

    Returns: (int) the stored random seed from the experiment
    """
    session = sessionmaker(bind=db_engine)()
    try:
        return session.query(TriageRun).get(run_id).random_seed
    finally:
        session.close()


def _write_csv(file_like, db_objects, type_of_object):
    writer = csv.writer(file_like, quoting=csv.QUOTE_MINIMAL, lineterminator="\n")
    for db_object in db_objects:
        if type(db_object) != type_of_object:
            raise TypeError(
                "Cannot copy collection of objects to db as they are not all "
                f"of the same type. First object was {type_of_object} "
                f"and later encountered a {type(db_object)}"
            )
        writer.writerow(
            [getattr(db_object, col.name) for col in db_object.__table__.columns]
        )


@db_retry
def save_db_objects(db_engine, db_objects):
    """Saves a collection of SQLAlchemy model objects to the database using a COPY command

    Args:
        db_engine (sqlalchemy.engine)
        db_objects (iterable) SQLAlchemy model objects, corresponding to a valid table
    """
    db_objects = iter(db_objects)
    first_object = next(db_objects)
    type_of_object = type(first_object)
    columns = [col.name for col in first_object.__table__.columns]

    with PipeTextIO(
        partial(
            _write_csv,
            db_objects=chain((first_object,), db_objects),
            type_of_object=type_of_object,
        )
    ) as pipe:
        postgres_copy.copy_from(
            source=pipe,
            dest=type_of_object,
            engine_or_conn=db_engine,
            columns=columns,
            format="csv",
        )
