import json
import time

from numpyencoder import NumpyEncoder
from omegaconf import OmegaConf
from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, relationship, sessionmaker

Base = declarative_base()


class Perflog(Base):
    __tablename__ = "perflog"
    id = Column(Integer, primary_key=True, index=True)
    dataset = Column(String)
    task = Column(String)
    model_type = Column(String)
    model_name = Column(String)
    hidden_dim = Column(Integer)
    performance = Column(String)
    config = Column(String)
    record_time = Column(Integer)


def process_performance_raw_info(
    cfg,
    mae=None,
    mse=None,
    rmse=None,
    mape=None,
    acc=None,
    auroc=None,
    auprc=None,
    early_prediction_score=None,
    multitask_prediction_score=None,
    verbose=0,
):
    result = []
    if mae is not None:
        result.extend(
            [
                {"name": "mae", "mean": mae.mean(), "std": mae.std()},
                {"name": "mse", "mean": mse.mean(), "std": mse.std()},
                {"name": "rmse", "mean": rmse.mean(), "std": rmse.std()},
                {"name": "mape", "mean": mape.mean(), "std": mape.std()},
            ]
        )
    if acc is not None:
        result.extend(
            [
                {"name": "acc", "mean": acc.mean(), "std": acc.std()},
                {"name": "auroc", "mean": auroc.mean(), "std": auroc.std()},
                {"name": "auprc", "mean": auprc.mean(), "std": auprc.std()},
            ]
        )
    thresholds = cfg.thresholds
    if early_prediction_score is not None:
        for i in range(len(thresholds)):
            result.append(
                {
                    "name": "early_prediction_score",
                    "mean": early_prediction_score.mean(axis=0)[i],
                    "std": early_prediction_score.std(axis=0)[i],
                    "threshold": thresholds[i],
                }
            )
    if multitask_prediction_score is not None:
        for i in range(len(thresholds)):
            result.append(
                {
                    "name": "multitask_prediction_score",
                    "mean": multitask_prediction_score.mean(axis=0)[i],
                    "std": multitask_prediction_score.std(axis=0)[i],
                    "threshold": thresholds[i],
                }
            )
    if verbose == 1:
        print(result)
    return result


def create_perflog(db: Session, cfg, perf=None):
    hid_dim = 0
    if "hidden_dim" in cfg:
        hid_dim = cfg.hidden_dim
    db_perflog = Perflog(
        task=cfg.task,
        dataset=cfg.dataset,
        model_type=cfg.model_type,
        model_name=cfg.model,
        hidden_dim=hid_dim,
        performance=json.dumps(perf, cls=NumpyEncoder),
        config=OmegaConf.to_yaml(cfg),
        record_time=int(time.time()),
    )
    db.add(db_perflog)
    db.commit()
    db.refresh(db_perflog)
    return db_perflog


def process_and_upload_performance(
    cfg,
    mae=None,
    mse=None,
    rmse=None,
    mape=None,
    acc=None,
    auroc=None,
    auprc=None,
    early_prediction_score=None,
    multitask_prediction_score=None,
    verbose=0,
    upload=False,
):
    perf = process_performance_raw_info(
        cfg,
        mae=mae,
        mse=mse,
        rmse=rmse,
        mape=mape,
        acc=acc,
        auroc=auroc,
        auprc=auprc,
        early_prediction_score=early_prediction_score,
        multitask_prediction_score=multitask_prediction_score,
        verbose=verbose,
    )
    if upload:
        db_cfg = OmegaConf.load("configs/_base_/db.yaml")

        engine, username, password, host, port, database = (
            db_cfg.engine,
            db_cfg.username,
            db_cfg.password,
            db_cfg.host,
            db_cfg.port,
            db_cfg.database,
        )
        SQLALCHEMY_DATABASE_URL = (
            f"{engine}://{username}:{password}@{host}:{port}/{database}"
        )
        engine = create_engine(SQLALCHEMY_DATABASE_URL)
        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

        db = SessionLocal()
        create_perflog(db=db, cfg=cfg, perf=perf)
        db.close()