from dataclasses import fields

import jax
import jax.numpy as jnp

import pgx
from pgx.utils import act_randomly


def validate(env: pgx.Env, num: int = 100):
    """validate checks these items:

    - init
      - reward is zero array
    - step
      - state.current_player is positive
      - (TODO) taking illegal actions terminates the episode with a negative reward
      - legal_action_mask is empty when terminated (TODO: or all True?)
      - taking actions at terminal states returns the same state (with zero reward)
      - (TODO) player change at the last step before terminal
    - observe
      - Returns different observations when player_ids are different (except the initial state)
    - TODO: reward must be zero when step is called after terminated
    - TODO: observation type (bool, int8 or int16) for efficiency; https://jax.readthedocs.io/en/latest/type_promotion.html
    """

    init = jax.jit(env.init)
    step = jax.jit(env.step)

    rng = jax.random.PRNGKey(849020)
    for _ in range(num):
        rng, subkey = jax.random.split(rng)
        state = init(subkey)

        assert state.steps == 0
        curr_steps = state.steps
        _validate_state(state)
        _validate_init_reward(state)
        _validate_current_player(state)
        _validate_legal_actions(state)

        while True:
            rng, subkey = jax.random.split(rng)
            action = act_randomly(subkey, state)
            state = step(state, action)
            assert (
                state.steps == curr_steps + 1
            ), f"{state.steps}, {curr_steps}"
            curr_steps += 1

            _validate_state(state)
            _validate_current_player(state)
            _validate_legal_actions(state)

            if state.terminated:
                break

        _validate_taking_action_after_terminal(state, step)


def _validate_taking_action_after_terminal(state: pgx.State, step_fn):
    prev_state = state
    if not state.terminated:
        return
    action = 0
    state = step_fn(state, action)
    assert (state.reward == 0).all()
    for field in fields(state):
        if field.name in ["reward", "steps"]:
            continue
        assert (
            getattr(state, field.name) == getattr(prev_state, field.name)
        ).all(), f"{field.name} : \n{getattr(state, field.name)}\n{getattr(prev_state, field.name)}"


def _validate_init_reward(state: pgx.State):
    assert (state.reward == jnp.zeros_like(state.reward)).all()


def _validate_state(state: pgx.State):
    """validate_state checks these items:

    - current_player is int8
    - terminated is bool_
    - reward is float
    - legal_action_mask is bool_
    - TODO: observation is bool_ or int8 (can promote to any other types)
    """
    assert state.current_player.dtype == jnp.int8, state.current_player.dtype
    assert state.terminated.dtype == jnp.bool_, state.terminated.dtype
    assert state.reward.dtype == jnp.float32, state.reward.dtype
    assert (
        state.legal_action_mask.dtype == jnp.bool_
    ), state.legal_action_mask.dtype


def _validate_legal_actions(state: pgx.State):
    if state.terminated:
        # Agent can take any action at terminal state (but give no effect to the next state)
        # This is to avoid zero-division error by normalizing action probability by legal actions
        assert (
            state.legal_action_mask == jnp.ones_like(state.legal_action_mask)
        ).all(), state.legal_action_mask
    else:
        ...


def _validate_current_player(state: pgx.State):
    assert (
        state.current_player >= 0
    ), f"current_player must be positive before terminated but got : {state.current_player}"
