import os
import typing
import json
import numpy as np
from markupsafe import Markup
from jinja2 import Template
from pyecharts.charts import Line, Bar, Page
from pyecharts import options as opts
from loguru import logger

from stagesepx.classifier import ClassifierResult, SingleClassifierResult
from stagesepx import toolbox
from stagesepx import constants
from stagesepx.cutter import VideoCutResult, VideoCutRange
from stagesepx.video import VideoFrame
from stagesepx import __VERSION__

# load template
template_dir_path = os.path.join(os.path.dirname(__file__), "template")
template_path = os.path.join(template_dir_path, "report.html")


def get_template() -> str:
    with open(template_path, encoding=constants.CHARSET) as t:
        template = t.read()
    return template


class Reporter(object):
    # 3 status:
    # - `stable` means nothing happened (nearly) during this period
    # - `unstable` means something happened
    # - `unspecific` means your model has no idea about `which class this frame should be` (lower than threshold)
    LABEL_STABLE: str = "stable"
    LABEL_UNSTABLE: str = "unstable"
    # unknown stage actually
    LABEL_UNSPECIFIC: str = "unspecific"

    def __init__(self):
        self.thumbnail_list: typing.List[typing.Tuple[str, str]] = list()
        self.extra_dict: typing.Dict[str, str] = dict()

    def add_thumbnail(self, name: str, pic_object: np.ndarray):
        """
        add picture object (cv2) to your report

        :param name:
        :param pic_object:
        :return:
        """
        b64_str = toolbox.np2b64str(pic_object)
        self.thumbnail_list.append((name, b64_str))

    def add_extra(self, name: str, value: str):
        """
        add some extra info ( key-value part) to your report

        :param name:
        :param value:
        :return:
        """
        self.extra_dict[name] = value

    @staticmethod
    def _draw_line(result: ClassifierResult) -> Line:
        # draw line chart
        x_axis = [str(i) for i in result.get_timestamp_list()]
        y_axis = result.get_stage_list()

        line = Line(init_opts=opts.InitOpts(bg_color=constants.BACKGROUND_COLOR))
        line.add_xaxis(x_axis)
        line.add_yaxis("stage", y_axis, is_step=False, is_symbol_show=True)
        line.set_global_opts(
            title_opts=opts.TitleOpts(
                title="Trend", subtitle="describe how these stages switching"
            ),
            toolbox_opts=opts.ToolboxOpts(is_show=True),
            tooltip_opts=opts.TooltipOpts(
                is_show=True, trigger="axis", axis_pointer_type="cross"
            ),
            brush_opts=opts.BrushOpts(x_axis_index="all", tool_box=["lineX"]),
        )
        return line

    @staticmethod
    def _draw_sim(data: VideoCutResult) -> Line:
        x_axis = [str(i.start) for i in data.range_list]
        ssim_axis = [i.ssim for i in data.range_list]
        mse_axis = [i.mse for i in data.range_list]
        psnr_axis = [i.psnr for i in data.range_list]

        line = Line(init_opts=opts.InitOpts(bg_color=constants.BACKGROUND_COLOR))
        line.add_xaxis(x_axis)
        line.add_yaxis("ssim", ssim_axis)
        line.add_yaxis("mse", mse_axis)
        line.add_yaxis("psnr", psnr_axis)
        line.set_global_opts(
            title_opts=opts.TitleOpts(title="SIM"),
            toolbox_opts=opts.ToolboxOpts(is_show=True),
            tooltip_opts=opts.TooltipOpts(
                is_show=True, trigger="axis", axis_pointer_type="cross"
            ),
        )
        line.set_series_opts(label_opts=opts.LabelOpts(is_show=False))
        return line

    @staticmethod
    def _draw_bar(result: ClassifierResult) -> Bar:
        # draw bar chart
        bar = Bar(init_opts=opts.InitOpts(bg_color=constants.BACKGROUND_COLOR))
        x_axis = sorted(list(result.get_stage_set()))
        y_axis = list()
        offset = result.get_offset()
        for each_stage_name in x_axis:
            ranges = result.get_specific_stage_range(each_stage_name)
            time_cost: float = 0.0
            for each in ranges:
                # last frame - first frame
                time_cost += each[-1].timestamp - each[0].timestamp + offset
            y_axis.append(time_cost)

        bar.add_xaxis(x_axis)
        bar.add_yaxis("time cost", y_axis)
        bar.set_global_opts(
            title_opts=opts.TitleOpts(title="Time Cost", subtitle="... of each stages"),
            toolbox_opts=opts.ToolboxOpts(is_show=True),
        )
        logger.debug(f"time cost: {dict(zip(x_axis, y_axis))}")
        return bar

    @staticmethod
    def get_stable_stage_sample_frame_list(
        result: ClassifierResult, *args, **kwargs
    ) -> typing.List[VideoFrame]:
        # VideoFrame: with data
        # SingleClassifierResult: without data
        last = result.data[0]
        picked: typing.List[SingleClassifierResult] = [last]
        for each in result.data:
            # ignore unstable stage
            if not each.is_stable():
                continue
            if last.stage != each.stage:
                last = each
                picked.append(each)

        return [each.to_video_frame(*args, **kwargs) for each in picked]

    @classmethod
    def get_stable_stage_sample(
        cls, result: ClassifierResult, *args, **kwargs
    ) -> np.ndarray:
        def get_split_line(f):
            return np.zeros((f.shape[0], 5))

        frame_list: typing.List[np.ndarray] = list()
        for each in cls.get_stable_stage_sample_frame_list(result, *args, **kwargs):
            frame_list.append(each.data)
            frame_list.append(get_split_line(each.data))
        return np.hstack(frame_list)

    @classmethod
    def save(cls, to_file: str, result: ClassifierResult):
        assert not os.path.isfile(to_file), f"file {to_file} already existed"
        data = [i.to_dict() for i in result.data]
        with open(to_file, "w", encoding=constants.CHARSET) as f:
            json.dump(data, f)

    @classmethod
    def load(cls, from_file: str) -> ClassifierResult:
        assert os.path.isfile(from_file), f"file {from_file} not existed"
        with open(from_file, encoding=constants.CHARSET) as f:
            content = json.load(f)
        return ClassifierResult([SingleClassifierResult(**each) for each in content])

    def draw(
        self,
        classifier_result: ClassifierResult,
        report_path: str = None,
        unstable_ranges: typing.List[VideoCutRange] = None,
        cut_result: VideoCutResult = None,
        compress_rate: float = None,
        target_size: typing.Tuple[int, int] = None,
        *_,
        **__,
    ):
        """
        draw report file

        :param classifier_result: classifierResult, output of classifier
        :param report_path: your report will be there
        :param unstable_ranges: for marking unstable ranges
        :param cut_result: more charts would be built
        :param compress_rate:
        :param target_size:
        :return:
        """
        # default: compress_rate
        if not compress_rate:
            compress_rate = 0.2
        if not unstable_ranges:
            unstable_ranges = []

        # draw
        line = self._draw_line(classifier_result)
        bar = self._draw_bar(classifier_result)

        # merge charts
        page = Page()
        page.add(line)
        page.add(bar)

        # insert pictures
        if cut_result:
            # sim chart
            sim_line = self._draw_sim(cut_result)
            page.add(sim_line)

        # mark range
        for each_range in unstable_ranges:
            classifier_result.mark_range_unstable(each_range.start, each_range.end)

        offset = classifier_result.get_offset()
        stage_range = classifier_result.get_stage_range()
        for cur_index in range(len(stage_range)):
            each_range = stage_range[cur_index]
            middle = each_range[len(each_range) // 2]
            # which means this range is stable
            if middle.is_stable():
                label = self.LABEL_STABLE
                frame = toolbox.compress_frame(
                    middle.get_data(),
                    compress_rate=compress_rate,
                    target_size=target_size,
                )
            # not stable
            else:
                # todo: looks not good enough. `unspecific` looks a little weird but I have no idea now
                if middle.stage == constants.UNKNOWN_STAGE_FLAG:
                    label = self.LABEL_UNSPECIFIC
                else:
                    label = self.LABEL_UNSTABLE
                # add a frame for human readable
                if cur_index + 1 < len(stage_range):
                    range_for_display = [*each_range, stage_range[cur_index + 1][0]]
                else:
                    range_for_display = each_range
                # merge these frames into one
                # note: these frames should have the same size
                frame = np.hstack(
                    [
                        toolbox.compress_frame(
                            i.get_data(),
                            compress_rate=compress_rate,
                            target_size=target_size,
                        )
                        for i in range_for_display
                    ]
                )

            first, last = each_range[0], each_range[-1]
            self.add_thumbnail(
                f"{label} range {first.frame_id}({first.timestamp}) - {last.frame_id}({last.timestamp + offset}), "
                f"duration: {last.timestamp - first.timestamp + offset}, "
                f"stage: {first.stage}",
                frame,
            )
        # calc time cost
        cost_dict = classifier_result.calc_changing_cost()

        # time stamp
        timestamp = toolbox.get_timestamp_str()

        # video
        self.add_extra("video path", classifier_result.video_path)
        self.add_extra("frame count", str(classifier_result.get_length()))
        self.add_extra("offset between frames", str(classifier_result.get_offset()))

        # insert extras
        template = Template(get_template())
        template_content = template.render(
            chart=Markup(page.render_embed()),
            thumbnail_list=self.thumbnail_list,
            extras=self.extra_dict,
            background_color=constants.BACKGROUND_COLOR,
            cost_dict=cost_dict,
            timestamp=timestamp,
            version_code=__VERSION__,
        )

        # default: write to current dir
        default_name = f"{timestamp}.html"
        if not report_path:
            report_path = default_name
        # somewhere specific
        # existed dir?
        elif os.path.isdir(report_path):
            report_path = os.path.join(report_path, default_name)
        logger.debug(f"trying to save report to {report_path}")

        # write file
        with open(report_path, "w", encoding=constants.CHARSET) as fh:
            fh.write(template_content)
        logger.info(f"save report to {report_path}")
