"""MinAtar/Seaquest: A fork of github.com/kenjyoung/MinAtar

The authors of original MinAtar implementation are:
    * Kenny Young (kjyoung@ualberta.ca)
    * Tian Tian (ttian@ualberta.ca)
The original MinAtar implementation is distributed under GNU General Public License v3.0
    * https://github.com/kenjyoung/MinAtar/blob/master/License.txt
"""
from typing import Literal, Optional

import jax
import jax.lax as lax
from jax import numpy as jnp

import pgx.v1 as v1
from pgx._src.struct import dataclass

RAMP_INTERVAL: jnp.ndarray = jnp.int32(100)
MAX_OXYGEN: jnp.ndarray = jnp.int32(200)
INIT_SPAWN_SPEED: jnp.ndarray = jnp.int32(20)
DIVER_SPAWN_SPEED: jnp.ndarray = jnp.int32(30)
INIT_MOVE_INTERVAL: jnp.ndarray = jnp.int32(5)
SHOT_COOL_DOWN: jnp.ndarray = jnp.int32(5)
ENEMY_SHOT_INTERVAL: jnp.ndarray = jnp.int32(10)
ENEMY_MOVE_INTERVAL: jnp.ndarray = jnp.int32(5)
DIVER_MOVE_INTERVAL: jnp.ndarray = jnp.int32(5)


ZERO: jnp.ndarray = jnp.int32(0)
NINE: jnp.ndarray = jnp.int32(9)
TRUE: jnp.ndarray = jnp.bool_(True)
FALSE: jnp.ndarray = jnp.bool_(False)


@dataclass
class State(v1.State):
    current_player: jnp.ndarray = jnp.int8(0)
    observation: jnp.ndarray = jnp.zeros((10, 10, 10), dtype=jnp.bool_)
    rewards: jnp.ndarray = jnp.zeros(1, dtype=jnp.float32)  # (1,)
    terminated: jnp.ndarray = FALSE
    truncated: jnp.ndarray = FALSE
    legal_action_mask: jnp.ndarray = jnp.ones(6, dtype=jnp.bool_)
    _rng_key: jax.random.KeyArray = jax.random.PRNGKey(0)
    _step_count: jnp.ndarray = jnp.int32(0)
    # --- MinAtar Seaquest specific ---
    _oxygen: jnp.ndarray = MAX_OXYGEN
    _diver_count: jnp.ndarray = ZERO
    _sub_x: jnp.ndarray = jnp.int32(5)
    _sub_y: jnp.ndarray = ZERO
    _sub_or: jnp.ndarray = FALSE
    _f_bullets: jnp.ndarray = -jnp.ones((5, 3), dtype=jnp.int32)
    _e_bullets: jnp.ndarray = -jnp.ones(
        (25, 3), dtype=jnp.int32
    )  # <= 1 per each sub
    _e_fish: jnp.ndarray = -jnp.ones((25, 4), dtype=jnp.int32)  # <= 19
    _e_subs: jnp.ndarray = -jnp.ones((25, 5), dtype=jnp.int32)  # <= 19
    _divers: jnp.ndarray = -jnp.ones((5, 4), dtype=jnp.int32)  # <= 2
    _e_spawn_speed: jnp.ndarray = INIT_SPAWN_SPEED
    _e_spawn_timer: jnp.ndarray = INIT_SPAWN_SPEED
    _d_spawn_timer: jnp.ndarray = DIVER_SPAWN_SPEED
    _move_speed: jnp.ndarray = INIT_MOVE_INTERVAL
    _ramp_index: jnp.ndarray = ZERO
    _shot_timer: jnp.ndarray = ZERO
    _surface: jnp.ndarray = TRUE
    _terminal: jnp.ndarray = FALSE
    _last_action: jnp.ndarray = ZERO

    @property
    def env_id(self) -> v1.EnvId:
        return "minatar-seaquest"

    def to_svg(
        self,
        *,
        color_theme: Optional[Literal["light", "dark"]] = None,
        scale: Optional[float] = None,
    ) -> str:
        del color_theme, scale
        from pgx.minatar.utils import visualize_minatar

        return visualize_minatar(self)

    def save_svg(
        self,
        filename,
        *,
        color_theme: Optional[Literal["light", "dark"]] = None,
        scale: Optional[float] = None,
    ) -> None:
        from pgx.minatar.utils import visualize_minatar

        visualize_minatar(self, filename)


class MinAtarSeaquest(v1.Env):
    def __init__(
        self,
        *,
        use_minimal_action_set: bool = True,
        sticky_action_prob: float = 0.1,
    ):
        super().__init__()
        self.use_minimal_action_set = use_minimal_action_set
        self.sticky_action_prob: float = sticky_action_prob
        self.minimal_action_set = jnp.int32([0, 1, 2, 3, 4, 5])
        self.legal_action_mask = jnp.ones(6, dtype=jnp.bool_)
        if self.use_minimal_action_set:
            self.legal_action_mask = jnp.ones(
                self.minimal_action_set.shape[0], dtype=jnp.bool_
            )

    def _init(self, key: jax.random.KeyArray) -> State:
        state = State()
        state = state.replace(legal_action_mask=self.legal_action_mask)  # type: ignore
        return state

    def _step(self, state: v1.State, action) -> State:
        assert isinstance(state, State)
        state = state.replace(legal_action_mask=self.legal_action_mask)  # type: ignore
        action = jax.lax.select(
            self.use_minimal_action_set,
            self.minimal_action_set[action],
            action,
        )
        return _step(state, action, sticky_action_prob=self.sticky_action_prob)  # type: ignore

    def _observe(self, state: v1.State, player_id: jnp.ndarray) -> jnp.ndarray:
        assert isinstance(state, State)
        return _observe(state)

    @property
    def id(self) -> v1.EnvId:
        return "minatar-seaquest"

    @property
    def version(self) -> str:
        return "beta"

    @property
    def num_players(self):
        return 1


def _step(
    state: State,
    action: jnp.ndarray,
    sticky_action_prob,
):
    key, subkey = jax.random.split(state._rng_key)
    state = state.replace(_rng_key=key)  # type: ignore
    rngs = jax.random.split(subkey, 6)
    action = jnp.int32(action)
    # sticky action
    action = jax.lax.cond(
        jax.random.uniform(rngs[0]) < sticky_action_prob,
        lambda: state._last_action,
        lambda: action,
    )
    enemy_lr = jax.random.choice(rngs[1], jnp.array([True, False]))
    is_sub = jax.random.choice(
        rngs[2], jnp.array([True, False]), p=jnp.array([1 / 3, 2 / 3])
    )
    enemy_y = jax.random.choice(rngs[3], jnp.arange(1, 9))
    diver_lr = jax.random.choice(rngs[4], jnp.array([True, False]))
    diver_y = jax.random.choice(rngs[5], jnp.arange(1, 9))
    return _step_det(
        state, action, enemy_lr, is_sub, enemy_y, diver_lr, diver_y
    )


def _step_det(
    state: State,
    action: jnp.ndarray,
    enemy_lr,
    is_sub,
    enemy_y,
    diver_lr,
    diver_y,
):
    ramping = TRUE

    oxygen = state._oxygen
    diver_count = state._diver_count
    sub_x = state._sub_x
    sub_y = state._sub_y
    sub_or = state._sub_or
    f_bullets = state._f_bullets
    e_bullets = state._e_bullets
    e_fish = state._e_fish
    e_subs = state._e_subs
    divers = state._divers
    e_spawn_speed = state._e_spawn_speed
    e_spawn_timer = state._e_spawn_timer
    d_spawn_timer = state._d_spawn_timer
    move_speed = state._move_speed
    ramp_index = state._ramp_index
    shot_timer = state._shot_timer
    surface = state._surface
    terminal = state._terminal

    r = jnp.float32(0)

    # Spawn enemy if timer is up
    e_subs, e_fish = lax.cond(
        e_spawn_timer == 0,
        lambda: _spawn_enemy(
            e_subs, e_fish, move_speed, enemy_lr, is_sub, enemy_y
        ),
        lambda: (e_subs, e_fish),
    )
    e_spawn_timer = lax.cond(
        e_spawn_timer == 0, lambda: e_spawn_speed, lambda: e_spawn_timer
    )

    # Spawn diver if timer is up
    divers, d_spawn_timer = lax.cond(
        d_spawn_timer == 0,
        lambda: (_spawn_diver(divers, diver_lr, diver_y), DIVER_SPAWN_SPEED),
        lambda: (divers, d_spawn_timer),
    )

    # Resolve player action
    f_bullets, shot_timer, sub_x, sub_y, sub_or = _resolve_action(
        action, shot_timer, f_bullets, sub_x, sub_y, sub_or
    )

    # Update friendly Bullets
    f_bullets, e_subs, e_fish, r = _update_friendly_bullets(
        f_bullets, e_subs, e_fish, r
    )

    # Update divers
    divers, diver_count = _update_divers(divers, diver_count, sub_x, sub_y)

    # Update enemy subs
    f_bullets, e_subs, e_bullets, terminal, r = _update_enemy_subs(
        f_bullets, e_subs, e_bullets, sub_x, sub_y, move_speed, terminal, r
    )

    # Update enemy bullets
    e_bullets, terminal = _update_enemy_bullets(
        e_bullets, sub_x, sub_y, terminal
    )

    # Update enemy fish
    f_bullets, e_fish, terminal, r = _update_enemy_fish(
        f_bullets, e_fish, sub_x, sub_y, move_speed, terminal, r
    )

    # Update various timers
    e_spawn_timer = lax.cond(
        e_spawn_timer > 0, lambda: e_spawn_timer - 1, lambda: e_spawn_timer
    )
    d_spawn_timer = lax.cond(
        d_spawn_timer, lambda: d_spawn_timer - 1, lambda: d_spawn_timer
    )
    shot_timer = lax.cond(
        shot_timer > 0, lambda: shot_timer - 1, lambda: shot_timer
    )
    terminal |= oxygen < 0
    tmp = surface
    oxygen = lax.cond(sub_y > 0, lambda: oxygen - 1, lambda: oxygen)
    surface = lax.cond(sub_y > 0, lambda: FALSE, lambda: surface)
    terminal = lax.cond(
        (sub_y <= 0) & ~tmp & (diver_count == 0),
        lambda: TRUE,
        lambda: terminal,
    )
    surface |= (sub_y <= 0) & ~tmp & (diver_count != 0)
    _r, oxygen, diver_count, move_speed, e_spawn_speed, ramp_index = lax.cond(
        (sub_y <= 0) & ~tmp & (diver_count != 0),
        lambda: _surface(
            diver_count, oxygen, e_spawn_speed, move_speed, ramping, ramp_index
        ),
        lambda: (
            jnp.int32(0),
            oxygen,
            diver_count,
            move_speed,
            e_spawn_speed,
            ramp_index,
        ),
    )
    r += _r

    state = state.replace(  # type: ignore
        _oxygen=oxygen,
        _diver_count=diver_count,
        _sub_x=sub_x,
        _sub_y=sub_y,
        _sub_or=sub_or,
        _f_bullets=f_bullets,
        _e_bullets=e_bullets,
        _e_fish=e_fish,
        _e_subs=e_subs,
        _divers=divers,
        _e_spawn_speed=e_spawn_speed,
        _e_spawn_timer=e_spawn_timer,
        _d_spawn_timer=d_spawn_timer,
        _move_speed=move_speed,
        _ramp_index=ramp_index,
        _shot_timer=shot_timer,
        _surface=surface,
        _terminal=terminal,
        _last_action=action,
        rewards=r[jnp.newaxis],
        terminated=terminal,
    )
    return state


def find_ix(arr):
    return (arr[:, 0] == -1).argmax()


def _resolve_action(action, shot_timer, f_bullets, sub_x, sub_y, sub_or):
    f_bullets, shot_timer = lax.cond(
        (action == 5) & (shot_timer == 0),
        lambda: (
            f_bullets.at[find_ix(f_bullets)].set(
                jnp.int32([sub_x, sub_y, sub_or])
            ),
            SHOT_COOL_DOWN,
        ),
        lambda: (f_bullets, shot_timer),
    )
    sub_x, sub_or = lax.cond(
        action == 1,
        lambda: (lax.max(ZERO, sub_x - 1), FALSE),
        lambda: (sub_x, sub_or),
    )
    sub_x, sub_or = lax.cond(
        action == 3,
        lambda: (lax.min(NINE, sub_x + 1), TRUE),
        lambda: (sub_x, sub_or),
    )
    sub_y = lax.cond(
        action == 2, lambda: lax.max(ZERO, sub_y - 1), lambda: sub_y
    )
    sub_y = lax.cond(
        action == 4, lambda: lax.min(jnp.int32(8), sub_y + 1), lambda: sub_y
    )
    return f_bullets, shot_timer, sub_x, sub_y, sub_or


def _update_by_f_bullets_hit(j, _f_bullets, e):
    is_hit = (_f_bullets[j, 0] == e[:, 0]) & (_f_bullets[j, 1] == e[:, 1])
    k = jnp.argmax(is_hit)
    k = jax.lax.select(jnp.sum(is_hit) == 0, 25, k)
    # k = lax.while_loop(
    #     lambda i: ~hit[i] & (i < 25),
    #     lambda i: i + 1,
    #     0,
    # )
    _f_bullets, e, removed = lax.cond(
        k < 25,
        lambda: (_remove_i(_f_bullets, j), _remove_i(e, k), TRUE),
        lambda: (_f_bullets, e, FALSE),
    )
    return _f_bullets, e, removed


def _update_friendly_bullets(f_bullets, e_subs, e_fish, r):
    def _remove(j, _f_bullets, _e_subs, _e_fish, _r):
        _f_bullets, _e_fish, fish_removed = _update_by_f_bullets_hit(
            j, _f_bullets, _e_fish
        )
        _r += fish_removed
        _f_bullets, _e_subs, sub_removed = lax.cond(
            fish_removed,
            lambda: (_f_bullets, _e_subs, FALSE),
            lambda: _update_by_f_bullets_hit(j, _f_bullets, _e_subs),
        )
        _r += sub_removed
        return _f_bullets, _e_subs, _e_fish, _r

    def _update_each(i, x):
        _f_bullets, _e_subs, _e_fish, _r = x
        j = 5 - i - 1
        is_filled = _is_filled(_f_bullets[j])
        _f_bullets = lax.cond(
            is_filled,
            lambda: _f_bullets.at[j, 0].add(
                lax.cond(_f_bullets[j, 2], lambda: 1, lambda: -1)
            ),
            lambda: _f_bullets,
        )
        _f_bullets, _e_subs, _e_fish, _r = lax.cond(
            is_filled,
            lambda: lax.cond(
                _is_out(_f_bullets[j]),
                lambda: (_remove_i(_f_bullets, j), _e_subs, _e_fish, _r),
                lambda: _remove(j, _f_bullets, _e_subs, _e_fish, _r),
            ),
            lambda: (_f_bullets, _e_subs, _e_fish, _r),
        )
        return _f_bullets, _e_subs, _e_fish, _r

    f_bullets, e_subs, e_fish, r = lax.fori_loop(
        0, 5, _update_each, (f_bullets, e_subs, e_fish, r)
    )
    return f_bullets, e_subs, e_fish, r


def _is_hit(row, x, y):
    return (row[0] == x) & (row[1] == y)


def _is_out(row):
    return (row[0] < 0) | (row[0] > 9)


def _is_filled(row):
    return jnp.any(row != -1)


def _update_divers(divers, diver_count, sub_x, sub_y):
    def _update_by_move(_divers, _diver_count, j):
        _divers = _divers.at[j, 3].set(DIVER_MOVE_INTERVAL)
        _divers = _divers.at[j, 0].add(
            lax.cond(_divers[j, 2], lambda: 1, lambda: -1)
        )
        _divers, _diver_count = lax.cond(
            _is_out(_divers[j]),
            lambda: (_remove_i(_divers, j), _diver_count),
            lambda: lax.cond(
                _is_hit(_divers[j], sub_x, sub_y) & (_diver_count < 6),
                lambda: (_remove_i(_divers, j), _diver_count + 1),
                lambda: (_divers, _diver_count),
            ),
        )
        return _divers, _diver_count

    def _update_each(i, x):
        _divers, _diver_count = x
        j = 5 - i - 1
        return lax.cond(
            _is_filled(_divers[j]),
            lambda: lax.cond(
                _is_hit(_divers[j], sub_x, sub_y) & (_diver_count < 6),
                lambda: (_remove_i(_divers, j), _diver_count + 1),
                lambda: lax.cond(
                    _divers[j, 3] == 0,
                    lambda: _update_by_move(_divers, _diver_count, j),
                    lambda: (_divers.at[j, 3].add(-1), _diver_count),
                ),
            ),
            lambda: (_divers, _diver_count),
        )

    divers, diver_count = lax.fori_loop(
        0, 5, _update_each, (divers, diver_count)
    )

    return divers, diver_count


def _update_enemy_subs(
    f_bullets, e_subs, e_bullets, sub_x, sub_y, move_speed, terminal, r
):
    def _update_sub(j, _f_bullets, _e_subs, _terminal, _r):
        _e_subs = _e_subs.at[j, 3].set(move_speed)
        _e_subs = _e_subs.at[j, 0].add(
            lax.cond(_e_subs[j, 2], lambda: 1, lambda: -1)
        )
        is_out = _is_out(_e_subs[j])
        is_hit = _is_hit(_e_subs[j], sub_x, sub_y)
        _e_subs = lax.cond(
            is_out, lambda: _remove_i(_e_subs, j), lambda: _e_subs
        )
        _terminal = lax.cond(~is_out & is_hit, lambda: TRUE, lambda: _terminal)
        _f_bullets, _e_subs, removed = lax.cond(
            ~is_out & ~is_hit,
            lambda: _update_by_hit(j, _f_bullets, _e_subs),
            lambda: (_f_bullets, _e_subs, FALSE),
        )
        _r += removed
        return (removed | is_out), _f_bullets, _e_subs, _terminal, _r

    def _update_each_filled(j, x):
        _f_bullets, _e_subs, _e_bullets, _terminal, _r = x
        _terminal |= _is_hit(_e_subs[j], sub_x, sub_y)
        removed, _f_bullets, _e_subs, _terminal, _r = lax.cond(
            _e_subs[j, 3] == 0,
            lambda: _update_sub(j, _f_bullets, _e_subs, _terminal, _r),
            lambda: (
                FALSE,
                _f_bullets,
                _e_subs.at[j, 3].add(-1),
                _terminal,
                _r,
            ),
        )
        timer_zero = _e_subs[j, 4] == 0
        _e_subs = lax.cond(
            removed,
            lambda: _e_subs,
            lambda: lax.cond(
                timer_zero,
                lambda: _e_subs.at[j, 4].set(ENEMY_SHOT_INTERVAL),
                lambda: _e_subs.at[j, 4].add(-1),
            ),
        )
        _e_bullets = lax.cond(
            removed,
            lambda: _e_bullets,
            lambda: lax.cond(
                timer_zero,
                lambda: _e_bullets.at[find_ix(_e_bullets)].set(
                    jnp.int32(
                        [
                            lax.cond(
                                _e_subs[j, 2],
                                lambda: _e_subs[j, 0],
                                lambda: _e_subs[j, 0],
                            ),
                            _e_subs[j, 1],
                            _e_subs[j, 2],
                        ]
                    )
                ),
                lambda: _e_bullets,
            ),
        )
        return _f_bullets, _e_subs, _e_bullets, _terminal, _r

    def _update_each(i, x):
        j = 25 - i - 1
        return lax.cond(
            _is_filled(x[1][j]), lambda: (_update_each_filled(j, x)), lambda: x
        )

    f_bullets, e_subs, e_bullets, terminal, r = lax.fori_loop(
        0, 25, _update_each, (f_bullets, e_subs, e_bullets, terminal, r)
    )

    return f_bullets, e_subs, e_bullets, terminal, r


def _remove_i(arr, i):
    mask = jnp.tile(jnp.arange(arr.shape[0]) < i, (arr.shape[1], 1)).T
    rolled = jnp.roll(arr, -1, axis=0)
    return jnp.where(mask, arr, rolled).at[-1, :].set(-1)


def _remove_out_of_bound(arr, ix):
    arr = lax.fori_loop(
        0,
        ix,
        lambda i, a: lax.cond(
            (a[i][0] < 0) & (a[i][1] > 9), lambda: _remove_i(a, i), lambda: a
        ),
        arr,
    )
    return arr


def _remove_hit(arr, ix, x, y):
    arr = lax.fori_loop(
        0,
        ix,
        lambda i, a: lax.cond(
            (a[i][0] == x) & (a[i][1] == y), lambda: _remove_i(a, i), lambda: a
        ),
        arr,
    )
    return arr


def _step_obj(arr, ix):
    arr_p = arr.at[:, 0].add(1)
    arr_m = arr.at[:, 0].add(-1)
    arr_2 = jnp.where(
        jnp.tile(arr[:, 2], reps=(arr.shape[1], 1)).T, arr_p, arr_m
    )
    arr = jnp.where(
        jnp.tile(jnp.arange(arr.shape[0]) < ix, reps=(arr.shape[1], 1)).T,
        arr_2,
        arr,
    )
    # arr = lax.fori_loop(
    #     0,
    #     ix,
    #     lambda i, a: a.at[i, 0].add(lax.cond(a[i, 2], lambda: 1, lambda: -1)),
    #     arr,
    # )
    return arr


def _hit(arr, ix, x, y):
    return (
        (arr[:, 0] == x) & (arr[:, 1] == y) & (jnp.arange(arr.shape[0]) < ix)
    ).any()


def _update_enemy_bullets(e_bullets, sub_x, sub_y, terminal):
    ix = find_ix(e_bullets)
    terminal |= _hit(e_bullets, ix, sub_x, sub_y)
    e_bullets = _step_obj(e_bullets, ix)
    e_bullets = _remove_out_of_bound(e_bullets, ix)
    terminal |= _hit(e_bullets, ix, sub_x, sub_y)
    return e_bullets, terminal


def _update_by_hit(j, _f_bullets, e):
    is_hit = (e[j, 0] == _f_bullets[:, 0]) & (e[j, 1] == _f_bullets[:, 1])
    k = jnp.argmax(is_hit)
    k = jax.lax.select(jnp.sum(is_hit) == 0, 5, k)
    # k = lax.while_loop(
    #     lambda i: ~_is_hit(e[j], _f_bullets[i, 0], _f_bullets[i, 1]) & (i < 5),
    #     lambda i: i + 1,
    #     0,
    # )
    _f_bullets, e, removed = lax.cond(
        k < 5,
        lambda: (_remove_i(_f_bullets, k), _remove_i(e, j), TRUE),
        lambda: (_f_bullets, e, FALSE),
    )
    return _f_bullets, e, removed


def _update_enemy_fish(
    f_bullets, e_fish, sub_x, sub_y, move_speed, terminal, r
):
    def _update_by_hit_fish(j, _f_bullets, e, _terminal, _r):
        _f_bullets, e, removed = _update_by_hit(j, _f_bullets, e)
        return _f_bullets, e, _terminal, _r + removed

    def _update_fish(j, _f_bullets, _e_fish, _terminal, _r):
        _e_fish = _e_fish.at[j, 3].set(move_speed)
        _e_fish = _e_fish.at[j, 0].add(
            lax.cond(_e_fish[j, 2], lambda: 1, lambda: -1)
        )
        _f_bullets, _e_fish, _terminal, _r = lax.cond(
            _is_out(_e_fish[j]),
            lambda: (_f_bullets, _remove_i(_e_fish, j), _terminal, _r),
            lambda: lax.cond(
                _is_hit(_e_fish[j], sub_x, sub_y),
                lambda: (_f_bullets, _e_fish, TRUE, _r),
                lambda: _update_by_hit_fish(
                    j, _f_bullets, _e_fish, _terminal, _r
                ),
            ),
        )
        return _f_bullets, _e_fish, _terminal, _r

    def _update_each(i, x):
        j = 25 - i - 1
        _f_bullets, _e_fish, _terminal, _r = x
        _terminal |= _is_hit(_e_fish[j], sub_x, sub_y)
        _f_bullets, _e_fish, _terminal, _r = lax.cond(
            _is_filled(_e_fish[j]),
            lambda: lax.cond(
                _e_fish[j, 3] == 0,
                lambda: _update_fish(j, _f_bullets, _e_fish, _terminal, _r),
                lambda: (_f_bullets, _e_fish.at[j, 3].add(-1), _terminal, _r),
            ),
            lambda: (_f_bullets, _e_fish, _terminal, _r),
        )
        return _f_bullets, _e_fish, _terminal, _r

    f_bullets, e_fish, terminal, r = lax.fori_loop(
        0, 25, _update_each, (f_bullets, e_fish, terminal, r)
    )

    return f_bullets, e_fish, terminal, r


# Called when player hits surface (top row) if they have no divers, this ends the game,
# if they have 6 divers this gives reward proportional to the remaining oxygen and restores full oxygen
# otherwise this reduces the number of divers and restores full oxygen
def _surface(
    diver_count, oxygen, e_spawn_speed, move_speed, ramping, ramp_index
):
    diver_count, r = lax.cond(
        diver_count == 6,
        lambda: (ZERO, oxygen * 10 // MAX_OXYGEN),
        lambda: (diver_count, jnp.int32(0)),
    )
    oxygen = MAX_OXYGEN
    diver_count -= 1
    ramp_update = ramping & ((e_spawn_speed > 1) | (move_speed > 2))
    ramp_index = lax.cond(
        ramp_update, lambda: ramp_index + 1, lambda: ramp_index
    )
    move_speed = lax.cond(
        ramp_update & ((move_speed > 2) & (ramp_index % 2)),
        lambda: move_speed - 1,
        lambda: move_speed,
    )
    e_spawn_speed = lax.cond(
        ramp_update & (e_spawn_speed > 1),
        lambda: e_spawn_speed - 1,
        lambda: e_spawn_speed,
    )
    return r, oxygen, diver_count, move_speed, e_spawn_speed, ramp_index


# Spawn an enemy fish or submarine in random row and random direction,
# if the resulting row and direction would lead to a collision, do nothing instead
def _spawn_enemy(e_subs, e_fish, move_speed, enemy_lr, is_sub, enemy_y):
    x = lax.cond(enemy_lr, lambda: ZERO, lambda: NINE)
    has_collision = (
        (e_subs[:, 1] == enemy_y) & (e_subs[:, 2] != enemy_lr)
    ).sum() > 0
    has_collision |= (
        (e_fish[:, 1] == enemy_y) & (e_fish[:, 2] != enemy_lr)
    ).sum() > 0
    return lax.cond(
        has_collision,
        lambda: (e_subs, e_fish),
        lambda: lax.cond(
            is_sub,
            lambda: (
                e_subs.at[find_ix(e_subs)].set(
                    jnp.int32(
                        [
                            x,
                            enemy_y,
                            enemy_lr,
                            move_speed,
                            ENEMY_SHOT_INTERVAL,
                        ]
                    )
                ),
                e_fish,
            ),
            lambda: (
                e_subs,
                e_fish.at[find_ix(e_fish)].set(
                    jnp.int32([x, enemy_y, enemy_lr, move_speed])
                ),
            ),
        ),
    )


# Spawn a diver in random row with random direction
def _spawn_diver(divers, diver_lr, diver_y):
    x = lax.cond(diver_lr, lambda: ZERO, lambda: NINE)
    ix = find_ix(divers)
    divers = divers.at[ix].set(
        jnp.array([x, diver_y, diver_lr, DIVER_MOVE_INTERVAL], dtype=jnp.int32)
    )
    return divers


def _observe(state: State) -> jnp.ndarray:
    obs = jnp.zeros((11, 11, 10), dtype=jnp.bool_)
    obs = obs.at[state._sub_y, state._sub_x, 0].set(TRUE)
    back_x = lax.cond(
        state._sub_or, lambda: state._sub_x - 1, lambda: state._sub_x + 1
    )
    obs = obs.at[state._sub_y, back_x, 1].set(TRUE)
    oxygen_guage = state._oxygen * 10 // MAX_OXYGEN
    # hotfix to align to original minatar
    oxygen_guage = lax.cond(
        state._oxygen < 0, lambda: jnp.int32(9), lambda: oxygen_guage
    )
    obs = obs.at[9, :, 7].set(
        jnp.where(jnp.arange(11) < oxygen_guage, TRUE, obs[9, :, 7])
    )
    mask = (9 - state._diver_count <= jnp.arange(11)) & (jnp.arange(11) < 9)
    obs = obs.at[9, :, 8].set(jnp.where(mask, TRUE, obs[9, :, 8]))
    obs = obs.at[state._f_bullets[:, 1], state._f_bullets[:, 0], 2].set(TRUE)
    obs = obs.at[state._e_bullets[:, 1], state._e_bullets[:, 0], 4].set(TRUE)

    obs = obs.at[state._e_fish[:, 1], state._e_fish[:, 0], 5].set(TRUE)
    back_x = (
        state._e_fish[:, 0]
        + jnp.array([1, -1], dtype=jnp.int32)[state._e_fish[:, 2]]
    )
    obs = obs.at[state._e_fish[:, 1], back_x, 3].set(TRUE)

    obs = obs.at[state._e_subs[:, 1], state._e_subs[:, 0], 6].set(TRUE)
    back_x = (
        state._e_subs[:, 0]
        + jnp.array([1, -1], dtype=jnp.int32)[state._e_subs[:, 2]]
    )
    obs = obs.at[state._e_subs[:, 1], back_x, 3].set(TRUE)

    obs = obs.at[state._divers[:, 1], state._divers[:, 0], 9].set(TRUE)
    back_x = (
        state._divers[:, 0]
        + jnp.array([1, -1], dtype=jnp.int32)[state._divers[:, 2]]
    )
    obs = obs.at[state._divers[:, 1], back_x, 3].set(TRUE)

    return obs[:10, :10, :]
