# -*- coding: utf-8 -*-
import random
import time

from typing import List, Dict, Tuple

from .shared import DreamImage


class RGBPalette:
    ID = "RGB_PALETTE"

    def __init__(self, colors: List[tuple[int, int, int]] = None, image: DreamImage = None):
        self._colors = []

        def _fix_tuple(t):
            if len(t) < 3:
                return (t[0], t[0], t[0])
            else:
                return t

        if image:
            for p, _, _ in image:
                self._colors.append(_fix_tuple(p))
        if colors:
            for c in colors:
                self._colors.append(_fix_tuple(c))

    def _calculate_channel_contrast(self, c):
        hist = list(map(lambda _: 0, range(16)))
        for pixel in self._colors:
            hist[pixel[c] // 16] += 1
        s = 0
        max_possible = (15 - 0) * (len(self) // 2) * (len(self) // 2)
        for i in range(16):
            for j in range(i):
                if i != j:
                    s += abs(i - j) * hist[i] * hist[j]
        return s / max_possible

    def _calculate_combined_contrast(self):
        s = 0
        for c in range(3):
            s += self._calculate_channel_contrast(c)
        return s / 3

    def analyze(self):
        total_red = 0
        total_blue = 0
        total_green = 0
        for pixel in self:
            total_red += pixel[0]
            total_green += pixel[1]
            total_blue += pixel[2]
        n = len(self._colors)
        r = float(total_red) / (255 * n)
        g = float(total_green) / (255 * n)
        b = float(total_blue) / (255 * n)
        return ((r + g + b) / 3.0, self._calculate_combined_contrast(), r, g, b)

    def __len__(self):
        return len(self._colors)

    def __iter__(self):
        return iter(self._colors)

    def random_iteration(self, seed=None):
        s = seed if seed is not None else int(time.time() * 1000)
        n = len(self._colors) - 1
        c = self._colors

        class _ColorIterator:
            def __init__(self):
                self._r = random.Random()
                self._r.seed(s)
                self._n = n
                self._c = c

            def __next__(self):
                return self._c[self._r.randint(0, self._n)]

        return _ColorIterator()


class PartialPrompt:
    ID = "PARTIAL_PROMPT"

    def __init__(self):
        self._data = {}

    def add(self, text: str, weight: float):
        output = PartialPrompt()
        output._data = dict(self._data)
        for parts in text.split(","):
            parts = parts.strip()
            if " " in parts:
                output._data["(" + parts + ")"] = weight
            else:
                output._data[parts] = weight
        return output

    def is_empty(self):
        return not self._data

    def abs_sum(self):
        if not self._data:
            return 0.0
        return sum(map(abs, self._data.values()))

    def abs_max(self):
        if not self._data:
            return 0.0
        return max(map(abs, self._data.values()))

    def scaled_by(self, f: float):
        new_data = PartialPrompt()
        new_data._data = dict(self._data)
        for text, weight in new_data._data.items():
            new_data._data[text] = weight * f
        return new_data

    def finalize(self, clamp: float):
        items = self._data.items()
        items = sorted(items, key=lambda pair: (pair[1], pair[0]))
        pos = list()
        neg = list()
        for text, w in sorted(items, key=lambda pair: (-pair[1], pair[0])):
            if w >= 0.0001:
                pos.append("({}:{:.3f})".format(text, min(clamp, w)))
        for text, w in sorted(items, key=lambda pair: (pair[1], pair[0])):
            if w <= -0.0001:
                neg.append("({}:{:.3f})".format(text, min(clamp, -w)))
        return ", ".join(pos), ", ".join(neg)


class LogEntry:
    ID = "LOG_ENTRY"

    @classmethod
    def new(cls, text):
        return LogEntry([(time.time(), text)])

    def __init__(self, data: List[Tuple[float, str]] = None):
        if data is None:
            self._data = list()
        else:
            self._data = list(data)

    def add(self, text: str):
        new_data = list(self._data)
        new_data.append((time.time(), text))
        return LogEntry(new_data)

    def merge(self, log_entry):
        new_data = list(self._data)
        new_data.extend(log_entry._data)
        return LogEntry(new_data)

    def get_filtered_entries(self, t: float):
        for d in sorted(self._data):
            if d[0] > t:
                yield d


class FrameCounter:
    ID = "FRAME_COUNTER"

    def __init__(self, current_frame=0, total_frames=1, frames_per_second=25.0):
        self.current_frame = max(0, current_frame)
        self.total_frames = max(total_frames, 1)
        self.frames_per_second = float(max(1.0, frames_per_second))

    def incremented(self, amount: int):
        return FrameCounter(self.current_frame + amount, self.total_frames, self.frames_per_second)

    @property
    def is_first_frame(self):
        return self.current_frame == 0

    @property
    def is_final_frame(self):
        return (self.current_frame + 1) == self.total_frames

    @property
    def is_after_last_frame(self):
        return self.current_frame >= self.total_frames

    @property
    def current_time_in_seconds(self):
        return float(self.current_frame) / self.frames_per_second

    @property
    def total_time_in_seconds(self):
        return float(self.total_frames) / self.frames_per_second

    @property
    def remaining_time_in_seconds(self):
        return self.total_time_in_seconds - self.current_time_in_seconds

    @property
    def progress(self):
        return float(self.current_frame) / (max(2, self.total_frames) - 1)


class AnimationSequence:
    ID = "ANIMATION_SEQUENCE"

    def __init__(self, frame_counter: FrameCounter, frames: Dict[int, List[str]] = None):
        self.frames = frames
        self.fps = frame_counter.frames_per_second
        self.frame_counter = frame_counter
        if self.is_defined:
            self.keys_in_order = sorted(frames.keys())
            self.num_batches = min(map(len, self.frames.values()))
        else:
            self.keys_in_order = []
            self.num_batches = 0

    @property
    def batches(self):
        return range(self.num_batches)

    def get_image_files_of_batch(self, batch_num):
        for key in self.keys_in_order:
            yield self.frames[key][batch_num]

    @property
    def is_defined(self):
        if self.frames:
            return True
        else:
            return False


class SharedTypes:
    frame_counter = {"frame_counter": (FrameCounter.ID,)}
    sequence = {"sequence": (AnimationSequence.ID,)}
    palette = {"palette": (RGBPalette.ID,)}