"""항구 입항 격자 세계: 강화학습 실습용 보조 모듈

배가 바람에 밀리는 바다에서 암초를 피해 항구로 들어가는 환경과
표 기반 Q 학습, 학습 과정 시각화 함수를 제공한다.
"""

from dataclasses import dataclass, field

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import animation, patches
from matplotlib.colors import LinearSegmentedColormap

# 행동 번호: FrozenLake와 같은 순서
LEFT, DOWN, RIGHT, UP = 0, 1, 2, 3
MOVES = {LEFT: (0, -1), DOWN: (1, 0), RIGHT: (0, 1), UP: (-1, 0)}

# S: 출항지, H: 항구, R: 암초, .: 바다
DEFAULT_MAP = [
    "S.......",
    "........",
    "R.RRRR..",
    "R.RRRR..",
    "........",
    "H.......",
]

SEA_CMAP = LinearSegmentedColormap.from_list(
    "sea", ["#f2f7fb", "#9ecae1", "#3182bd", "#08306b"]
)


class HarborEnv(gym.Env):
    """암초를 피해 항구로 들어가는 격자 세계

    - wind: 바람에 밀릴 확률. 의도한 방향으로 1 - wind,
      좌우(진행 방향의 수직)로 각각 wind / 2만큼 밀림
    - 항구 도착: 보상 +1, 종료
    - 암초 충돌: 보상 0, 종료
    - step_cost: 매 이동마다 빼는 연료비
    """

    metadata = {"render_modes": []}

    def __init__(self, wind=0.0, step_cost=0.0, desc=None, max_steps=100):
        super().__init__()
        self.desc = np.array([list(row) for row in (desc or DEFAULT_MAP)])
        self.nrow, self.ncol = self.desc.shape
        self.wind = wind
        self.step_cost = step_cost
        self.max_steps = max_steps
        self.observation_space = gym.spaces.Discrete(self.nrow * self.ncol)
        self.action_space = gym.spaces.Discrete(4)
        self.start = self.to_state(*np.argwhere(self.desc == "S")[0])
        self.P = self._build_transitions()
        self.s = self.start
        self.t = 0

    # 좌표와 상태 번호 변환
    def to_state(self, row, col):
        return int(row * self.ncol + col)

    def to_rc(self, s):
        return divmod(int(s), self.ncol)

    def cell(self, s):
        r, c = self.to_rc(s)
        return self.desc[r, c]

    def is_terminal(self, s):
        return self.cell(s) in "HR"

    def _move(self, s, a):
        r, c = self.to_rc(s)
        dr, dc = MOVES[a]
        r = min(max(r + dr, 0), self.nrow - 1)  # 벽에 부딪히면 제자리
        c = min(max(c + dc, 0), self.ncol - 1)
        return self.to_state(r, c)

    def _outcomes(self, a):
        """행동 a를 했을 때 실제 이동 방향과 확률"""
        if self.wind == 0:
            return [(1.0, a)]
        side = self.wind / 2
        return [(1 - self.wind, a), (side, (a - 1) % 4), (side, (a + 1) % 4)]

    def _build_transitions(self):
        """P[s][a] = [(확률, 다음 상태, 보상, 종료 여부), ...]"""
        P = {}
        for s in range(self.observation_space.n):
            P[s] = {}
            for a in range(4):
                P[s][a] = []
                if self.is_terminal(s):
                    P[s][a].append((1.0, s, 0.0, True))
                    continue
                for p, actual in self._outcomes(a):
                    s2 = self._move(s, actual)
                    reward = (1.0 if self.cell(s2) == "H" else 0.0) - self.step_cost
                    P[s][a].append((p, s2, reward, self.is_terminal(s2)))
        return P

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.s = self.start
        self.t = 0
        return self.s, {}

    def step(self, action):
        outcomes = self._outcomes(int(action))
        probs = [p for p, _ in outcomes]
        actual = outcomes[self.np_random.choice(len(outcomes), p=probs)][1]
        s2 = self._move(self.s, actual)
        reward = (1.0 if self.cell(s2) == "H" else 0.0) - self.step_cost
        self.s = s2
        self.t += 1
        terminated = self.is_terminal(s2)
        truncated = self.t >= self.max_steps
        return s2, reward, terminated, truncated, {"drift": actual != action}


# ---------------------------------------------------------------------------
# 학습
# ---------------------------------------------------------------------------


@dataclass
class History:
    """학습 기록"""

    Q: np.ndarray
    snapshots: list = field(default_factory=list)  # (에피소드, Q 복사본)
    success: list = field(default_factory=list)  # 에피소드별 입항 성공 여부

    @property
    def V(self):
        return self.Q.max(axis=1)


def greedy(Q, s, rng):
    """가치가 가장 큰 행동 선택(같으면 무작위)"""
    best = np.flatnonzero(Q[s] == Q[s].max())
    return int(rng.choice(best))


def q_learning(
    env,
    episodes=1000,
    alpha=0.1,
    gamma=0.99,
    epsilon=0.1,
    q_init=0.0,
    snapshot_every=10,
    seed=0,
):
    """표 기반 Q 학습

    Q(s, a) ← Q(s, a) + α [r + γ max Q(s', ·) − Q(s, a)]

    - q_init: 행동 가치의 초기값. 1로 두면 가 보지 않은 곳을
      좋을 것이라고 낙관하고 시도해 보게 됨(낙관적 초기값)
    """
    rng = np.random.default_rng(seed)
    Q = np.full((env.observation_space.n, 4), float(q_init))
    for s in range(env.observation_space.n):
        if env.is_terminal(s):
            Q[s] = 0.0  # 항구와 암초는 이후 보상이 없음
    history = History(Q=Q, snapshots=[(0, Q.copy())])
    env.reset(seed=seed)
    for ep in range(1, episodes + 1):
        s, _ = env.reset()
        done = False
        while not done:
            if rng.random() < epsilon:
                a = int(rng.integers(4))  # 탐색: 무작위 행동
            else:
                a = greedy(Q, s, rng)  # 활용: 가치가 가장 큰 행동
            s2, r, terminated, truncated, _ = env.step(a)
            target = r if terminated else r + gamma * Q[s2].max()
            Q[s, a] += alpha * (target - Q[s, a])
            s = s2
            done = terminated or truncated
        history.success.append(env.cell(s) == "H")
        if ep % snapshot_every == 0:
            history.snapshots.append((ep, Q.copy()))
    return history


def value_iteration(env, gamma=0.99, tol=1e-8):
    """전이 확률을 알 때 최적 행동 가치(정답) 계산"""
    nS = env.observation_space.n
    V = np.zeros(nS)
    while True:
        Q = np.zeros((nS, 4))
        for s in range(nS):
            for a in range(4):
                Q[s, a] = sum(
                    p * (r + (0 if done else gamma * V[s2]))
                    for p, s2, r, done in env.P[s][a]
                )
        V_new = Q.max(axis=1)
        if np.abs(V_new - V).max() < tol:
            return Q
        V = V_new


def run_episode(env, Q=None, seed=None):
    """한 번 항해하고 경로를 반환(Q가 없으면 무작위 행동)"""
    rng = np.random.default_rng(seed)
    s, _ = env.reset(seed=seed)
    path, drifts = [s], []
    done = False
    while not done:
        a = int(rng.integers(4)) if Q is None else greedy(Q, s, rng)
        s, _, terminated, truncated, info = env.step(a)
        path.append(s)
        drifts.append(info["drift"])
        done = terminated or truncated
    return path, drifts


def success_rate(env, Q=None, n=1000, seed=0):
    """n번 항해했을 때 입항에 성공한 비율"""
    ok = 0
    for i in range(n):
        path, _ = run_episode(env, Q, seed=seed + i)
        ok += env.cell(path[-1]) == "H"
    return ok / n


# ---------------------------------------------------------------------------
# 그리기
# ---------------------------------------------------------------------------


def _draw_rock(ax, r, c):
    pts = np.array([[-0.38, 0.2], [-0.2, -0.32], [0.15, -0.38], [0.4, -0.05],
                    [0.3, 0.33], [-0.1, 0.4]])
    ax.add_patch(patches.Polygon(pts + [c, r], closed=True,
                                 fc="#5b4636", ec="#2d2219", lw=1.5, zorder=3))


def _draw_port(ax, r, c):
    ax.add_patch(patches.Rectangle((c - 0.45, r - 0.45), 0.9, 0.9,
                                   fc="#ffe9a8", ec="#b8860b", lw=2, zorder=3))
    ax.text(c, r, "PORT", ha="center", va="center", fontsize=9,
            fontweight="bold", color="#7a5300", zorder=4)


def _draw_ship(ax, r, c, color="#d62728", size=1.0):
    hull = np.array([[-0.3, -0.05], [0.3, -0.05], [0.2, 0.18], [-0.2, 0.18]])
    cabin = np.array([[-0.08, -0.05], [0.12, -0.05], [0.12, -0.22], [-0.08, -0.22]])
    for pts, fc in [(hull, color), (cabin, "white")]:
        ax.add_patch(patches.Polygon(pts * size + [c, r], closed=True,
                                     fc=fc, ec="black", lw=1, zorder=5))


def _draw_board(ax, env, values=None, vmax=1.0, show_numbers=True):
    nrow, ncol = env.nrow, env.ncol
    if values is None:
        values = np.zeros(nrow * ncol)
    grid = values.reshape(nrow, ncol).copy()
    ax.imshow(grid, cmap=SEA_CMAP, vmin=0, vmax=vmax)
    for s in range(nrow * ncol):
        r, c = env.to_rc(s)
        cell = env.cell(s)
        if cell == "R":
            ax.add_patch(patches.Rectangle((c - 0.5, r - 0.5), 1, 1,
                                           fc="#cfc6b8", zorder=2))
            _draw_rock(ax, r, c)
        elif cell == "H":
            _draw_port(ax, r, c)
        elif show_numbers:
            v = grid[r, c]
            ax.text(c + 0.42, r + 0.42, f"{v:.2f}", ha="right", va="bottom",
                    fontsize=7, color="white" if v > vmax * 0.6 else "#333")
    ax.set_xticks(np.arange(-0.5, ncol, 1), minor=True)
    ax.set_yticks(np.arange(-0.5, nrow, 1), minor=True)
    ax.grid(which="minor", color="white", lw=1)
    ax.tick_params(which="both", length=0, labelbottom=False, labelleft=False)


def _draw_policy(ax, env, Q):
    for s in range(env.observation_space.n):
        if env.is_terminal(s) or np.allclose(Q[s], Q[s, 0]):
            continue  # 아직 차이가 없는 칸은 화살표 생략
        r, c = env.to_rc(s)
        dr, dc = MOVES[int(Q[s].argmax())]
        ax.arrow(c - dc * 0.18, r - dr * 0.18, dc * 0.3, dr * 0.3,
                 head_width=0.15, head_length=0.1, fc="#ff7f0e", ec="#ff7f0e",
                 lw=2, zorder=4, length_includes_head=False)


def plot_map(env, ax=None):
    """바다 지도 그리기"""
    own = ax is None  # 새 그림을 만들었으면 끝에서 화면에 출력
    if own:
        _, ax = plt.subplots(figsize=(env.ncol * 0.8, env.nrow * 0.8))
    _draw_board(ax, env, show_numbers=False)
    _draw_ship(ax, *env.to_rc(env.start))
    ax.set_title(f"Harbor (wind={env.wind})")
    if own:
        plt.show()


def plot_values(env, Q, ax=None, title=None, arrows=True):
    """칸별 가치(색)와 가장 좋은 행동(화살표)"""
    own = ax is None  # 새 그림을 만들었으면 끝에서 화면에 출력
    if own:
        _, ax = plt.subplots(figsize=(env.ncol * 0.8, env.nrow * 0.8))
    _draw_board(ax, env, Q.max(axis=1))
    _draw_ship(ax, *env.to_rc(env.start), size=0.6)
    if arrows:
        _draw_policy(ax, env, Q)
    if title:
        ax.set_title(title)
    if own:
        plt.show()


def animate(env, history, interval=100, save=None):
    """학습 과정 애니메이션: 기록된 시점(기본 10 에피소드)마다 한 장씩 같은 간격으로 재생

    - 노트북에서 슬라이더로 재생, save에 파일명을 주면 GIF 저장
    """
    from IPython.display import HTML

    idx = range(len(history.snapshots))
    fig, ax = plt.subplots(figsize=(env.ncol * 0.8, env.nrow * 0.8), dpi=80)

    def draw(i):
        ax.clear()
        ep, Q = history.snapshots[i]
        plot_values(env, Q, ax=ax, title=f"episode {ep}")

    anim = animation.FuncAnimation(fig, draw, frames=idx, interval=interval)
    if save:
        anim.save(save, writer="pillow", fps=1000 // interval)
    plt.close(fig)
    return HTML(anim.to_jshtml())


def plot_routes(env, Q=None, n=5, seed=0, ax=None, title=None):
    """n번 항해한 경로 그리기(Q가 없으면 무작위 항해)

    - 실선: 경로, ×: 바람에 밀린 지점
    - 끝점: 초록 원은 입항, 빨간 원은 좌초 또는 시간 초과
    """
    own = ax is None  # 새 그림을 만들었으면 끝에서 화면에 출력
    if own:
        _, ax = plt.subplots(figsize=(env.ncol * 0.8, env.nrow * 0.8))
    _draw_board(ax, env, show_numbers=False)
    _draw_ship(ax, *env.to_rc(env.start))
    colors = plt.cm.tab10.colors
    arrived = 0
    for i in range(n):
        path, drifts = run_episode(env, Q, seed=seed + i)
        rc = np.array([env.to_rc(s) for s in path], dtype=float)
        rc += (i - (n - 1) / 2) * 0.06  # 겹치지 않게 살짝 이동
        color = colors[i % 10]
        ax.plot(rc[:, 1], rc[:, 0], "-", color=color, lw=2, alpha=0.8, zorder=6)
        d = np.flatnonzero(drifts) + 1
        ax.plot(rc[d, 1], rc[d, 0], "x", color=color, ms=7, mew=2, zorder=7)
        ok = env.cell(path[-1]) == "H"
        arrived += ok
        ax.plot(rc[-1, 1], rc[-1, 0], "o", ms=10, zorder=8,
                mfc="#2ca02c" if ok else "#d62728", mec="white", mew=1.5)
    ax.set_title(title or f"{'greedy' if Q is not None else 'random'} routes: "
                 f"{arrived}/{n} arrived")
    if own:
        plt.show()


def plot_curve(history, window=100, ax=None):
    """입항 성공률의 이동평균"""
    own = ax is None  # 새 그림을 만들었으면 끝에서 화면에 출력
    if own:
        _, ax = plt.subplots(figsize=(6, 3))
    s = np.array(history.success, dtype=float)
    avg = np.convolve(s, np.ones(window) / window, mode="valid")
    ax.plot(np.arange(window, len(s) + 1), avg)
    ax.set_ylim(0, 1.02)
    ax.set_xlabel("episode")
    ax.set_ylabel(f"success rate ({window} ep avg)")
    ax.grid(alpha=0.3)
    if own:
        plt.show()