from __future__ import annotations

import numbers
import numpy as np
import sys
from itertools import accumulate

from ..Qt import QtWidgets, QtGui, QtCore
from .._py.gol import Ui_Form


def circle(radius):
    pxl = []
    for x in range(int(np.ceil(radius/np.sqrt(2)))):
        y = round(np.sqrt(radius**2-x**2))
        pxl.extend([[x, y], [y, x], [x, -y], [y, -x],
                    [-x, -y], [-y, -x], [-x, y], [-y, x]])

    return np.array(pxl)


def square(a):
    pxl = []
    pxl.extend(list(zip(range(-a, a+1), [a]*(2*a+1))))
    pxl.extend(list(zip(range(-a, a+1), [-a]*(2*a+1))))
    pxl.extend(list(zip([a]*(2*a+1), range(-a, a+1))))
    pxl.extend(list(zip([-a]*(2*a+1), range(-a, a+1))))

    return np.array(pxl)


def diamond(a):
    pxl = []
    for x in range(int(a+1)):
        y = a-x
        pxl.extend([[x, y], [-x, y], [x, -y], [-x, -y]])

    return np.array(pxl)


def plus(a):
    pxl = np.zeros((4*int(a)+2, 2), dtype=int)
    pxl[:2*int(a)+1, 0] = np.arange(-a, a+1)
    pxl[2*int(a)+1:, 1] = np.arange(-a, a+1)

    return pxl

# birth, survival
predefined_rules = {
    'Conway': ('23', '3'),
    'Life34': ('34', '34'),
    'Coagulation': ('235678', '378'),
    'Corrosion': ('12345', '45'),
    'Long life': ('5', '345'),
    'Maze': ('12345', '3'),
    'Coral': ('45678', '3'),
    'Pseudo life': ('238', '357'),
    'Flakes': ('012345678', '3'),
    'Gnarl': ('1', '1'),
    'Fabric': ('12', '1'),
    'Assimilation': ('4567', '345'),
    'Diamoeba': ('5678', '35678'),
    'High life': ('23', '36'),
    'More Maze': ('1235', '3'),
    'Replicator': ('1357', '1357'),
    'Seed': ('', '2'),
    'Serviette': ('', '234'),
    'More coagulation': ('235678', '3678'),
    'Domino': ('125', '36'),
    'Anneal': ('35678', '4678'),
}


class GameOfLife:
    colors = [
        [31, 119, 180],
        [255, 127, 14],
        [44, 160, 44],
        [214, 39, 40],
        [148, 103, 189],
        [140, 86, 75]
    ]

    def __init__(self,
                 size: tuple = (400, 400),
                 pattern: np.ndarray = None,
                 fill: float | list[float] = 0.05,
                 num_pop: int = 1,
                 rules: str | tuple = ('23', '3')
                 ):
        self.populations = num_pop if pattern is None else 1
        self._size = size

        self._world = np.zeros(shape=(*self._size, self.populations), dtype=np.uint8)
        self._neighbors = np.zeros(shape=self._world.shape, dtype=np.uint8)
        self._drawing = 255 * np.zeros(shape=(*self._size, 3), dtype=np.uint8)

        self.fill = np.zeros(self.populations)
        self._populate(fill, pattern)

        if isinstance(rules, str):
            try:
                b_rule, s_rule = predefined_rules[rules]
            except KeyError:
                raise ValueError('Rule is not predefined')
        else:
            b_rule, s_rule = rules

        self.survival_condition = np.array([int(c) for c in s_rule])
        self.birth_condition = np.array([int(c) for c in b_rule])

        # indexes for neighbors
        self._neighbor_idx = [
            [[slice(None), slice(1, None)], [slice(None), slice(0, -1)]],      # N (:, 1:), S (:, _-1)
            [[slice(1, None), slice(1, None)], [slice(0, -1), slice(0, -1)]],  # NE (1:, 1:), SW (:-1, :-1)
            [[slice(1, None), slice(None)], [slice(0, -1), slice(None)]],      # E (1:, :), W (:-1, :)
            [[slice(1, None), slice(0, -1)], [slice(0, -1), slice(1, None)]]   # SE (1:, :-1), NW (:-1:, 1:)
        ]

    def _populate(self, fill, pattern):
        if pattern is None:
            if isinstance(fill, numbers.Number):
                fill = [fill]*self.populations

            prob = (np.random.default_rng().random(size=self._size))
            lower_lim = 0
            for i, upper_lim in enumerate(accumulate(fill)):
                self._world[:, :, i] = (lower_lim <= prob) & (prob < upper_lim)
                lower_lim = upper_lim

        else:
            pattern = np.asarray(pattern)

            x_step = self._size[0]//(self.populations+1)
            y_step = self._size[1]//(self.populations+1)

            for i in range(self.populations):
                self._world[-pattern[:, 0]+(i+1)*x_step, pattern[:, 1]+(i+1)*y_step, i] = 1

        for i in range(self.populations):
            self.fill[i] = self._world[:, :, i].sum() / (self._size[0]*self._size[1])

    def tick(self):
        n = self._neighbors
        w = self._world

        n[...] = 0
        for idx_1, idx_2 in self._neighbor_idx:
            n[tuple(idx_1)] += w[tuple(idx_2)]
            n[tuple(idx_2)] += w[tuple(idx_1)]

        birth = ((np.in1d(n, self.birth_condition).reshape(n.shape)) & (w.sum(axis=-1) == 0)[:, :, None])
        survive = ((np.in1d(n, self.survival_condition).reshape(n.shape)) & (w == 1))

        w[...] = 0
        w[birth | survive] = 1

    def draw(self, shade: int):
        if shade == 0:
            self._drawing[...] = 0
        elif shade == 1:
            self._drawing -= (self._drawing/4).astype(np.uint8)
            self._drawing = self._drawing.clip(0, 127)

        for i in range(self.populations):
            self._drawing[(self._world[:, :, i] == 1)] = self.colors[i]
            self.fill[i] = self._world[:, :, i].sum() / (self._size[0]*self._size[1])

        return self._drawing


class QGameOfLife(QtWidgets.QDialog, Ui_Form):
    SPEEDS = [0.5, 1, 2, 5, 7.5, 10, 12.5, 15, 20, 25, 30, 40, 50]

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        self._size = (100, 100)
        self.game = None
        self._step = 0
        self._shading = 1
        self._speed = 5

        self.timer = QtCore.QTimer()
        self.timer.setInterval(100)
        self.timer.timeout.connect(self.tick)

        self._init_ui()

    def _init_ui(self):
        self.item = None
        self.scene = QtWidgets.QGraphicsScene()
        self.item = self.scene.addPixmap(QtGui.QPixmap())
        self.view.setScene(self.scene)

        self.rule_cb.addItems(list(predefined_rules.keys()))

        self.random_widgets = []
        for _ in range(6):
            w = QSliderText(15, parent=self)
            w.slider.valueChanged.connect(self.set_max_population)
            self.verticalLayout.addWidget(w)
            self.random_widgets.append(w)

        self.birth_line.setValidator(QtGui.QIntValidator())
        self.survival_line.setValidator(QtGui.QIntValidator())

        for w in self.random_widgets[1:] + [self.object_widget]:
            w.hide()

        self.setGeometry(QtWidgets.QStyle.alignedRect(QtCore.Qt.LeftToRight, QtCore.Qt.AlignCenter,
                                                      self.size(), QtWidgets.qApp.desktop().availableGeometry()))

        self.view.resizeEvent = self.resizeEvent

    @QtCore.pyqtSlot(int)
    def on_object_combobox_currentIndexChanged(self, idx: int):
        for w in self.random_widgets + [self.object_widget, self.rand_button_wdgt]:
            w.hide()

        if idx == 0:
            self.random_widgets[0].show()
            self.rand_button_wdgt.show()
        else:
            self.object_widget.show()

    @QtCore.pyqtSlot()
    def on_add_random_button_clicked(self):
        if self.object_combobox.currentIndex() != 0:
            return

        for w in self.random_widgets[1:]:
            if not w.isVisible():
                w.show()
                break

    @QtCore.pyqtSlot()
    def on_remove_random_button_clicked(self):
        if self.object_combobox.currentIndex() != 0:
            return

        for w in reversed(self.random_widgets[1:]):
            if w.isVisible():
                w.hide()
                break

    @QtCore.pyqtSlot(str)
    def on_rule_cb_currentIndexChanged(self, entry: str):
        rule = predefined_rules[entry]
        self.birth_line.setText(rule[1])
        self.survival_line.setText(rule[0])

    @QtCore.pyqtSlot(int)
    def set_max_population(self, _: int):
        over_population = -100
        num_tribes = -1
        for w in self.random_widgets:
            if w.isVisible():
                over_population += w.slider.value()
                num_tribes += 1

        if over_population > 0:
            for w in self.random_widgets:
                if w == self.sender() or w.isHidden():
                    continue
                w.setValue(max(0, int(w.slider.value()-over_population/num_tribes)))

    @QtCore.pyqtSlot()
    def on_start_button_clicked(self):
        self.pause_button.setChecked(False)
        self._step = 0
        self.current_step.setText(f'{self._step} steps')
        self._size = (int(self.height_box.value()), int(self.width_box.value()))

        pattern = None
        num_pop = 0
        fill = []
        pattern_size = self.object_size.value()
        if 2*pattern_size >= max(self._size):
            pattern_size = int(np.floor(max(self._size[0]-1, self._size[1]-1) / 2))

        idx = self.object_combobox.currentIndex()
        if idx == 0:
            for w in self.random_widgets:
                if w.isVisible():
                    num_pop += 1
                    fill.append(w.slider.value()/100)
        else:
            pattern = [None, circle, square, diamond, plus][idx](pattern_size)

        self.game = GameOfLife(self._size, pattern=pattern, fill=fill, num_pop=num_pop,
                               rules=(self.birth_line.text(), self.survival_line.text()))
        self.draw()
        self.view.fitInView(self.item, QtCore.Qt.KeepAspectRatio)

        self.timer.start()

    def tick(self):
        self.game.tick()
        self.draw()

        self._step += 1
        self.current_step.setText(f'{self._step} steps')
        self.cover_label.setText('\n'.join([f'Color {i+1}: {f*100:.2f} %' for i, f in enumerate(self.game.fill)]))

    @QtCore.pyqtSlot()
    def on_faster_button_clicked(self):
        self._speed = min(self._speed+1, len(QGameOfLife.SPEEDS)-1)
        new_speed = QGameOfLife.SPEEDS[self._speed]
        self.timer.setInterval(int(1000/new_speed))
        self.velocity_label.setText(f'{new_speed:.1f} steps/s')

    @QtCore.pyqtSlot()
    def on_slower_button_clicked(self):
        self._speed = max(self._speed-1, 0)
        new_speed = QGameOfLife.SPEEDS[self._speed]
        self.timer.setInterval(int(1000/new_speed))
        self.velocity_label.setText(f'{new_speed:.1f} steps/s')

    @QtCore.pyqtSlot()
    def on_pause_button_clicked(self):
        if self.pause_button.isChecked():
            self.timer.stop()
        else:
            self.timer.start()

    @QtCore.pyqtSlot(QtWidgets.QAbstractButton)
    def on_buttonGroup_buttonClicked(self, button):
        self._shading = [self.radioButton, self.vanish_shadow, self.full_shadow].index(button)

    def draw(self):
        bitmap = self.game.draw(shade=self._shading)
        h, w, c = bitmap.shape
        image = QtGui.QImage(bitmap.tobytes(), w, h, w*c, QtGui.QImage.Format_RGB888)
        self.scene.removeItem(self.item)
        pixmap = QtGui.QPixmap.fromImage(image)
        self.item = self.scene.addPixmap(pixmap)

    @QtCore.pyqtSlot(int)
    def on_hide_button_stateChanged(self, state: int):
        self.option_frame.setVisible(not state)
        self.view.fitInView(self.scene.sceneRect(), QtCore.Qt.KeepAspectRatio)

    def resizeEvent(self, evt):
        super().resizeEvent(evt)
        self.view.fitInView(self.item, QtCore.Qt.KeepAspectRatio)

    def showEvent(self, evt):
        super().showEvent(evt)
        self.view.fitInView(self.scene.sceneRect(), QtCore.Qt.KeepAspectRatio)


class QSliderText(QtWidgets.QWidget):
    def __init__(self, value: int, parent=None):
        super().__init__(parent=parent)
        layout = QtWidgets.QHBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(3)

        self.slider = QtWidgets.QSlider(self)
        self.slider.setOrientation(QtCore.Qt.Horizontal)
        self.slider.setMaximum(100)
        self.slider.setTickPosition(self.slider.TicksBothSides)
        self.slider.setTickInterval(5)

        self.value = QtWidgets.QLabel(self)

        self.slider.valueChanged.connect(lambda x: self.value.setText(f'{x} %'))

        layout.addWidget(self.slider)
        layout.addWidget(self.value)

        self.setLayout(layout)

        self.setValue(value)

    def setValue(self, value: int):
        self.slider.setValue(value)
        self.value.setText(f'{value} %')
