"""
Ising Model for the MCH2022 Badge

Implements the [[https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm][Metropolis-Hastings algorithm]].

And probably in a bad way.

See Also
--------
[[https://en.wikipedia.org/wiki/Square_lattice_Ising_model]]:
    Background info on the Ising model.
"""
import math
import display
import urandom
import buttons
import mch22
import framebuf

# the initial temperature
tau = 0.4
# Critical temperature for 2d Ising
# this is 2 / Arsinh(1)
tau_curie = 2.269185314213022

# Note: max size is 320×240
size = (64, 48)

# FrameBuffer needs 2 bytes for every RGB565 pixel
spins = framebuf.FrameBuffer(
    bytearray(size[0] * size[1] // 8),
    size[0], size[1], framebuf.MONO_HMSB)


# scale up
size_pixel = (display.width() // size[0],
              display.height() // size[1])
border = ((display.width() - size_pixel[0] * size[0]) // 2,
          (display.height() - size_pixel[1] * size[1]) // 2)
display.translate(border[0], border[1])

color_down = 0x000000
color_up = 0xffffff


def get_spin_value(i, j):
    """
    return the value of the spin at (i, j)

    If i, j are too large or negative, use periodic boundary conditions.
    """
    return 2 * spins.pixel(
        (i + size[0]) % size[0],
        (j + size[1]) % size[1]) - 1


def set_spin_value(i, j, value):
    """
    set the value of a spin
    """
    spins.pixel(i, j, (value + 1) // 2)
    color = get_color_from_value(value)
    display.drawRect(
        i * size_pixel[0], j * size_pixel[1],
        size_pixel[0], size_pixel[1],
        True, color)


def reboot(pressed):
    """
    reboot on button press
    """
    if pressed:
        mch22.exit_python()


def create_random_state():
    """
    start with a random configuration
    """
    # Okay, we start with the most stupid setup first:
    for i in range(size[0]):
        for j in range(size[1]):
            set_spin_value(i, j, 2 * urandom.randint(0, 1) - 1)

    display_info(0)
    display.flush()


def display_info(step):
    """
    Print info about temperature on screen
    """
    m = get_magnetization()
    text_color = 0xff0000
    font = "roboto_regular12"
    display.drawRect(5, 5, 297, 25, True, 0x9f9f9f)
    display.drawText(
        10, 10,
        f"Ising Magnet ({size[0]}x{size[1]}) , T = {tau:.1f}, M = {m:.3f} (step {step})",
        text_color, font)


def get_magnetization():
    """
    calculate the average magnetization
    """
    m = 0
    for i in range(size[0]):
        for j in range(size[1]):
            m += get_spin_value(i, j)
    return m / (size[0] * size[1])


def get_color_from_value(value):
    """
    The inverse of the above
    """
    value = max(-1, min(1, value))
    if value > 0.0:
        return color_up
    else:
        return color_down


def metropolis():
    """
    Perform n·m Metropolis steps
    """
    for _ in range(size[0] * size[1]):
        # pick a random pixel
        i = urandom.randint(0, size[0] - 1)
        j = urandom.randint(0, size[1] - 1)

        ij = get_spin_value(i, j)
        ij_left = get_spin_value(i - 1, j)
        ij_right = get_spin_value(i + 1, j)
        ij_up = get_spin_value(i, j - 1)
        ij_down = get_spin_value(i, j + 1)
        # difference in energy if spin is flipped
        deltaE = 2 * ij * (ij_left + ij_right + ij_up + ij_down)

        if deltaE < 0 or (
            # if energy is lower - flip spin
            # else only flip with exponentially small probability
                urandom.uniform(0, 1) < math.exp(- deltaE / (tau * tau_curie))):
            new_value = -ij
            set_spin_value(i, j, new_value)

    global step
    display_info(step)
    display.flush()


def increase_temp(pressed):
    """
    increase the temperature by 0.1
    """
    if not pressed:
        return
    global tau
    tau += 0.1
    display_info(step)


def decrease_temp(pressed):
    """
    decrease temperature by 0.1
    """
    if not pressed:
        return
    global tau
    tau -= 0.1
    # make sure we are not running into negative or zero temperatures
    tau = max(0.1, tau)
    display_info(step)


buttons.attach(buttons.BTN_MENU, reboot)
buttons.attach(buttons.BTN_UP, increase_temp)
buttons.attach(buttons.BTN_DOWN, decrease_temp)

# start with a grey screen as background
color_bg = 0x7f7f7f
display.drawFill(color_bg)

create_random_state()
step = 0
while True:
    step += 1
    metropolis()