import json

import os

import os.path as osp

from datetime import datetime

import numpy as np

import plotly.graph_objects as go

import torch

import wfdb

from tqdm import tqdm

from utils.network_utils import load_checkpoint

class BasePipeline:

    def __init__(self, config):

        self.config = config

        self.exp_name = self.config.get("exp_name", None)

        if self.exp_name is None:

            self.exp_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

        self.res_dir = osp.join(self.config["exp_dir"], self.exp_name, "results")

        os.makedirs(self.res_dir, exist_ok=True)

        self.model = self._init_net()

        self.pipeline_loader = self._init_dataloader()

        self.mapper = json.load(open(config["mapping_json"]))

        self.mapper = {j: i for i, j in self.mapper.items()}

        pretrained_path = self.config.get("model_path", False)

        if pretrained_path:

            load_checkpoint(pretrained_path, self.model)

        else:

            raise Exception(

                "model_path doesnt't exist in config. Please specify checkpoint path",

            )

    def _init_net(self):

        raise NotImplemented

    def _init_dataloader(self):

        raise NotImplemented

    def run_pipeline(self):

        self.model.eval()

        pd_class = np.empty(0)

        pd_peaks = np.empty(0)

        with torch.no_grad():

            for i, batch in tqdm(enumerate(self.pipeline_loader)):

                inputs = batch["image"].to(self.config["device"])

                predictions = self.model(inputs)

                classes = predictions.topk(k=1)[1].view(-1).cpu().numpy()

                pd_class = np.concatenate((pd_class, classes))

                pd_peaks = np.concatenate((pd_peaks, batch["peak"]))

        pd_class = pd_class.astype(int)

        pd_peaks = pd_peaks.astype(int)

        annotations = []

        for label, peak in zip(pd_class, pd_peaks):

            if (

                peak < len(self.pipeline_loader.dataset.signal)

                and self.mapper[label] != "N"

            ):

                annotations.append(

                    {

                        "x": peak,

                        "y": self.pipeline_loader.dataset.signal[peak],

                        "text": self.mapper[label],

                        "xref": "x",

                        "yref": "y",

                        "showarrow": True,

                        "arrowcolor": "black",

                        "arrowhead": 1,

                        "arrowsize": 2,

                    },

                )

        if osp.exists(self.config["ecg_data"] + ".atr"):

            ann = wfdb.rdann(self.config["ecg_data"], extension="atr")

            for label, peak in zip(ann.symbol, ann.sample):

                if peak < len(self.pipeline_loader.dataset.signal) and label != "N":

                    annotations.append(

                        {

                            "x": peak,

                            "y": self.pipeline_loader.dataset.signal[peak] - 0.1,

                            "text": label,

                            "xref": "x",

                            "yref": "y",

                            "showarrow": False,

                            "bordercolor": "#c7c7c7",

                            "borderwidth": 1,

                            "borderpad": 4,

                            "bgcolor": "#ffffff",

                            "opacity": 1,

                        },

                    )

        fig = go.Figure(

            data=go.Scatter(

                x=list(range(len(self.pipeline_loader.dataset.signal))),

                y=self.pipeline_loader.dataset.signal,

            ),

        )

        fig.update_layout(

            title="ECG",

            xaxis_title="Time",

            yaxis_title="ECG Output Value",

            title_x=0.5,

            annotations=annotations,

            autosize=True,

        )

        fig.write_html(

            osp.join(self.res_dir, osp.basename(self.config["ecg_data"] + ".html")),

        )