import hsi
from hsi import preprocessing, util
import numpy as np
import requests
import io
import base64
import matplotlib.pyplot as plt
from matplotlib import colors, colormaps
import asyncio
import os
import pandas as pd
import ast

from PIL import Image

import altair as alt

from pyscript import when, display, document, fetch, fs


def my_exec(script, globals=None, locals=None):
    '''Execute a script and return the value of the last expression'''
    stmts = list(ast.iter_child_nodes(ast.parse(script)))
    if not stmts:
        return None
    if isinstance(stmts[-1], ast.Expr):
        # the last one is an expression and we will try to return the results
        # so we first execute the previous statements
        if len(stmts) > 1:
            exec(compile(ast.Module(body=stmts[:-1]), filename="<ast>", mode="exec"), globals, locals)
        # then we eval the last one
        return eval(compile(ast.Expression(body=stmts[-1].value), filename="<ast>", mode="eval"), globals, locals)
    else:
        # otherwise we just execute the entire code
        return exec(script, globals, locals)


himg_out = None


async def handle_event(event):
    print(event.code)
    global himg_out
    res = my_exec(event.code)
    if isinstance(res, hsi.HSImage):
        himg_out = res
        await update_fig(None)
    return False


editor = document.querySelector("#editor")
editor.handleEvent = handle_event


async def init():
    global himg_out
    himg_out = my_exec(editor.innerHTML)
    await update_fig(None)


asyncio.create_task(init())


def new_fig(channel):
    axis = get_axis()
    a = himg_out.array_plane(channel, axis)

    norm = colors.Normalize(vmin=a.min(), vmax=a.max())
    cmap = colormaps.get_cmap('magma')
    rgba = cmap(norm(a))
    rgb = (rgba[:, :, :3] * 255).astype(np.uint8)

    img = Image.fromarray(rgb)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    img_str = 'data:image/png;base64,' + base64.b64encode(buf.read()).decode('UTF-8')
    return img_str


def get_axis():
    axis = document.querySelector("#axis").selectedIndex
    match axis:
        case 0:
            axis = hsi.bands
        case 1:
            axis = hsi.lines
        case 2:
            axis = hsi.samples
        case _:
            raise ValueError("Invalid axis selected")
    return axis


@when('change', '#axis')
async def update_axis(event):
    await update_fig(None)


@when('input', '#channel')
async def update_fig(event):
    axis = document.querySelector("#axis").selectedIndex
    match axis:
        case 0:
            sh = himg_out.shape.bands
        case 1:
            sh = himg_out.shape.lines
        case 2:
            sh = himg_out.shape.samples
        case _:
            raise ValueError("Invalid axis selected")

    document.querySelector("#channel").max = sh - 1
    val = int(document.querySelector("#channel").value)
    document.querySelector("#figure").src = new_fig(val)


@when('mousemove', '#figure')
def handle_fig_click(event):
    x, y = event.pageX, event.pageY
    left = document.querySelector("#figure").offsetLeft
    top = document.querySelector("#figure").offsetTop
    x -= left
    y -= top

    a = himg_out[y - 10:y + 10, x - 10:x + 10, :].ensure_dtype(hsi.float32).mean_axis(hsi.lines).mean_axis(
        hsi.samples).to_numpy().flatten()

    df = pd.DataFrame({
        "Band": np.arange(a.size),
        "Intensity": a
    })

    chart = alt.Chart(df).mark_line().encode(
        x="Band",
        y="Intensity"
    )

    display(chart, target="#chart", append=False)