# -- Code Cell --
import numpy as np
import pandas as pd

# ══════════════════════════════════════════════════════════════════════════════
# CONFIG
# ══════════════════════════════════════════════════════════════════════════════
GRID            = 32
N_STATES        = GRID * GRID       # 1024
N_ACTIONS       = 4
GOAL_STATE      = 777      # 1023 → (31,31)
GOAL_REWARD     = 1000.0
STEP_REWARD     = -1.0
GAMMA           = 0.95
EPSILON         = 1e-4
RANDOM_SEED     = 42
N_SAMPLES_TRAIN = 50_000
TRAIN_SPLIT     = 0.8               # 80% of states → train, 20% → test
ACTION_NAMES    = ["up", "down", "left", "right"]

# slot 0 = intended, 1 = perp-left, 2 = perp-right, 3 = backwards
OUTCOME_DIRS = {
    0: [0, 2, 3, 1],
    1: [1, 3, 2, 0],
    2: [2, 0, 1, 3],
    3: [3, 1, 0, 2],
}
DELTAS = {0: (-1, 0), 1: (1, 0), 2: (0, -1), 3: (0, 1)}

def idx(r, c):
    return max(0, min(GRID - 1, r)) * GRID + max(0, min(GRID - 1, c))


# ══════════════════════════════════════════════════════════════════════════════
# STEP 1 — Split states into train / test (disjoint, stratified by zone)
# ══════════════════════════════════════════════════════════════════════════════
print("=" * 60)
print("STEP 1 — Splitting states into train / test (disjoint)")
print("=" * 60)

rng_split = np.random.default_rng(RANDOM_SEED)

def get_zone(s):
    row, col = divmod(s, GRID)
    if row < GRID // 2 and col < GRID // 2:
        return "reliable"
    elif row >= GRID // 2 and col >= GRID // 2:
        return "slippery"
    else:
        return "mixed"

zones = {"reliable": [], "slippery": [], "mixed": []}
for s in range(N_STATES):
    zones[get_zone(s)].append(s)

train_states_set = set()
test_states_set  = set()

for zone_name, zone_states in zones.items():
    arr  = np.array(zone_states)
    rng_split.shuffle(arr)
    cut  = int(len(arr) * TRAIN_SPLIT)
    train_states_set.update(arr[:cut].tolist())
    test_states_set.update(arr[cut:].tolist())

# Goal state must be in train so value iteration sees it
train_states_set.add(GOAL_STATE)
test_states_set.discard(GOAL_STATE)

train_states = np.array(sorted(train_states_set))
test_states  = np.array(sorted(test_states_set))

assert len(set(train_states) & set(test_states)) == 0, "OVERLAP DETECTED!"

print(f"  Total states  : {N_STATES}")
print(f"  Train states  : {len(train_states)}  ({len(train_states)/N_STATES*100:.0f}%)")
print(f"  Test  states  : {len(test_states)}   ({len(test_states)/N_STATES*100:.0f}%)")
print(f"  Overlap check : NONE (disjoint confirmed)")

for zone_name in ["reliable", "slippery", "mixed"]:
    zs   = set(zones[zone_name])
    n_tr = len(zs & set(train_states.tolist()))
    n_te = len(zs & set(test_states.tolist()))
    print(f"    {zone_name:8s}: {n_tr} train  {n_te} test")


# ══════════════════════════════════════════════════════════════════════════════
# STEP 2 — Build per-state transition probabilities
# ══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("STEP 2 — Building per-state transition probabilities")
print("=" * 60)

def build_state_probs(seed=RANDOM_SEED):
    rng         = np.random.default_rng(seed)
    state_probs = np.zeros((N_STATES, N_ACTIONS, 4))

    for s in range(N_STATES):
        for a in range(N_ACTIONS):
            alpha = [
                rng.uniform(4, 6),     # intended  — moderate forward bias
                rng.uniform(0.8, 1.2), # perp-left
                rng.uniform(0.8, 1.2), # perp-right
                rng.uniform(0.3, 0.7), # backwards
            ]
            state_probs[s, a, :] = rng.dirichlet(alpha)

    return state_probs

state_probs = build_state_probs(seed=RANDOM_SEED)
print(f"  state_probs shape: {state_probs.shape}  (states × actions × outcome-slots)")

for zone_name in ["reliable", "slippery", "mixed"]:
    zs    = list(zones[zone_name])
    p_int = state_probs[zs, :, 0].mean()
    print(f"  Avg P(intended) [{zone_name:8s}]: {p_int:.3f}")


# ══════════════════════════════════════════════════════════════════════════════
# STEP 3 — Generate train.csv
#           Columns: state, action, next_state, reward
#           Only train_states are used as starting states → no test leakage
# ══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("STEP 3 — Generating train.csv (train states only)")
print("=" * 60)

def move(state, action, rng):
    row, col      = divmod(state, GRID)
    outcome_k     = rng.choice(4, p=state_probs[state, action])
    actual_action = OUTCOME_DIRS[action][outcome_k]
    dr, dc        = DELTAS[actual_action]
    return idx(row + dr, col + dc)

def sample_transitions(source_states, n_samples, seed):
    rng  = np.random.default_rng(seed)
    rows = []

    # --- Phase 1: cover every (state, action) pair at least once ---
    for s in source_states:
        for a in range(N_ACTIONS):
            ns = move(int(s), int(a), rng)
            r  = GOAL_REWARD if ns == GOAL_STATE else STEP_REWARD
            rows.append((int(s), int(a), ns, r))

    # --- Phase 2: random samples to reach n_samples total ---
    remaining = max(0, n_samples - len(rows))
    if remaining > 0:
        idxs    = rng.integers(0, len(source_states), size=remaining)
        states  = source_states[idxs]
        actions = rng.integers(0, N_ACTIONS, size=remaining)
        for s, a in zip(states, actions):
            ns = move(int(s), int(a), rng)
            r  = GOAL_REWARD if ns == GOAL_STATE else STEP_REWARD
            rows.append((int(s), int(a), ns, r))

    return pd.DataFrame(rows, columns=["state", "action", "next_state", "reward"])

df_train = sample_transitions(train_states, N_SAMPLES_TRAIN, seed=RANDOM_SEED)

# Hard check: no test states in the 'state' column
leaked = set(df_train["state"].tolist()) & set(test_states.tolist())
assert len(leaked) == 0, f"LEAK! {leaked}"

# train.csv → only the 4 columns a model needs to learn P and R
df_train[["state", "action", "next_state", "reward"]].to_csv("train.csv", index=False)

print(f"  train.csv  → {len(df_train):,} rows")
print(f"  Columns    : state, action, next_state, reward")
print(f"  Goal hits  : {(df_train.next_state == GOAL_STATE).sum()}")
print(f"  Unique train states covered: {df_train.state.nunique()} / {len(train_states)}")
print(f"  Train/test state leak check: NONE")
print("\n  Sample rows:")
print(df_train[["state", "action", "next_state", "reward"]].head(6).to_string(index=False))


# ══════════════════════════════════════════════════════════════════════════════
# STEP 4 — Build P and R from df_train
# ══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("STEP 4 — Building P and R matrices from train.csv")
print("=" * 60)

count_sa  = np.zeros((N_STATES, N_ACTIONS))
count_sas = np.zeros((N_STATES, N_ACTIONS, N_STATES))
reward_sa = np.zeros((N_STATES, N_ACTIONS))

for row in df_train.itertuples(index=False):
    s, a, ns, r = int(row.state), int(row.action), int(row.next_state), float(row.reward)
    count_sa[s, a]      += 1
    count_sas[s, a, ns] += 1
    reward_sa[s, a]     += r

seen = count_sa > 0
P = np.zeros((N_STATES, N_ACTIONS, N_STATES))
for s in range(N_STATES):
    for a in range(N_ACTIONS):
        if seen[s, a]:
            P[s, a, :] = count_sas[s, a, :] / count_sa[s, a]
        else:
            # Unseen (s,a): use the known per-state transition probs to build P exactly.
            # This avoids both the uniform-teleport bug (inflated Q) and stay-in-place bias.
            row_s, col_s = divmod(s, GRID)
            for k in range(4):
                actual_a = OUTCOME_DIRS[a][k]
                dr, dc   = DELTAS[actual_a]
                ns_      = idx(row_s + dr, col_s + dc)
                P[s, a, ns_] += state_probs[s, a, k]

R = np.where(seen, reward_sa / np.where(seen, count_sa, 1), STEP_REWARD)

print(f"  P shape : {P.shape}")
print(f"  R shape : {R.shape}")
print(f"  (s,a) pairs seen: {int(seen.sum()):,} / {N_STATES * N_ACTIONS:,}")
print(f"  P sums to 1 for all seen (s,a): {np.allclose(P[seen].sum(axis=1), 1.0)}")


# ══════════════════════════════════════════════════════════════════════════════
# STEP 5 — Value iteration → V* and Q*
# ══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("STEP 5 — Value iteration")
print("=" * 60)

V = np.zeros(N_STATES)
Q = np.zeros((N_STATES, N_ACTIONS))

for iteration in range(1, 10_001):
    Q                = R + GAMMA * np.einsum("ijk,k->ij", P, V)
    Q[GOAL_STATE, :] = 0.0
    V_new            = Q.max(axis=1)
    delta            = np.abs(V_new - V).max()
    V                = V_new

    if iteration % 50 == 0 or delta < EPSILON:
        print(f"  iter {iteration:5d}  |ΔV|_max = {delta:.8f}")

    if delta < EPSILON:
        print(f"\n  Converged at iteration {iteration}!")
        break

policy = Q.argmax(axis=1)
print(f"\n  V*(state=0)    = {V[0]:.4f}")
print(f"  V*(state=1023) = {V[GOAL_STATE]:.4f}   (terminal)")


# ══════════════════════════════════════════════════════════════════════════════
# STEP 6 — test.csv  → only 'state' column (unseen states, no labels)
#           ground_truth.csv → state + optimal action (for F1 eval)
# ══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("STEP 6 — Generating test.csv and ground_truth.csv")
print("=" * 60)

optimal_actions = Q[test_states, :].argmax(axis=1)

# test.csv — what you hand to the model: just the state
pd.DataFrame({"state": test_states}).to_csv("test.csv", index=False)

# ground_truth.csv — kept separate, used only for evaluation
pd.DataFrame({
    "state" : test_states,
    "action": optimal_actions,          # integer label (0-3)
}).to_csv("ground_truth.csv", index=False)

print(f"  test.csv         → {len(test_states):,} rows  |  columns: state")
print(f"  ground_truth.csv → {len(test_states):,} rows  |  columns: state, action")

print(f"\n  Zone distribution in test set:")
for zone_name in ["reliable", "slippery", "mixed"]:
    zs = set(zones[zone_name])
    n  = sum(s in zs for s in test_states)
    print(f"    {zone_name:8s}: {n:3d} states  ({n/len(test_states)*100:.1f}%)")

print(f"\n  Optimal action distribution (ground truth):")
for a in range(N_ACTIONS):
    n = (optimal_actions == a).sum()
    print(f"    {ACTION_NAMES[a]:5s}: {n:3d} states  ({n/len(test_states)*100:.1f}%)")


# ══════════════════════════════════════════════════════════════════════════════
# STEP 7 — Example evaluation using ground_truth.csv  (F1 score)
# ══════════════════════════════════════════════════════════════════════════════
print("\n" + "=" * 60)
print("STEP 7 — Example evaluation (F1 score)")
print("=" * 60)

from sklearn.metrics import f1_score, classification_report

# Simulate: load test.csv, predict argmax Q for each state, compare to ground truth
df_test_in = pd.read_csv("test.csv")
df_gt      = pd.read_csv("ground_truth.csv")

y_pred = Q[df_test_in["state"].values, :].argmax(axis=1)
y_true = df_gt["action"].values

macro_f1 = f1_score(y_true, y_pred, average="macro")
print(f"\n  Macro F1 : {macro_f1:.4f}")
print(f"\n  Per-class report:")
print(classification_report(y_true, y_pred, labels=list(range(N_ACTIONS)), target_names=ACTION_NAMES))

print("\nDone.  Files written: train.csv, test.csv, ground_truth.csv")

# -- Code Cell --


# -- Code Cell --


# -- Code Cell --
