"""Suit-invariant public features; privileged labels live outside this module."""
import numpy as np
import torch

from qojbot import rules as R
from qojbot.adapter import latest_log

KINDS = ('pass', 'single', 'pair', 'trio', 'trio1', 'trio2', 'bomb',
         'rocket', 'four2', 'four22', 'straight', 'pairs', 'plane',
         'plane1', 'plane2', 'bid')
ROLES = ('landlord', 'landlord_down', 'landlord_up')
HISTORY = 96
ACTION_DIM = 30 + len(KINDS) + 4
HISTORY_DIM = 15 + len(KINDS) + 3 + 4
SCALAR_DIM = 15


def role_seats(landlord):
    return (landlord, (landlord + 1) % 3, (landlord + 2) % 3)


def action_features(hand, moves):
    result = np.zeros((len(moves), ACTION_DIM), dtype=np.float32)
    own = np.asarray(R.counts(hand))
    for index, move in enumerate(moves):
        cards = np.asarray(R.counts(move.code))
        result[index, :15] = cards / 4
        result[index, 15:30] = (own - cards) / 4
        result[index, 30 + KINDS.index(move.kind)] = 1
        result[index, -4:] = ((move.rank + 1) / 15, move.length / 12,
                              move.size / 20, float(move.size == R.size(hand)))
    return torch.from_numpy(result)


def encode_public(obs, state, moves):
    """Only Observation and actor's state are accepted, never engine hands."""
    seats = role_seats(obs.landlord)
    role = seats.index(obs.seat)
    played = np.zeros((3, 15), dtype=np.float32)
    history = []
    for event in latest_log(state):
        kind = event.get('kind')
        if kind not in ('play', 'pass', 'bid'):
            continue
        row = np.zeros(HISTORY_DIM, dtype=np.float32)
        actor = seats.index(event['seat'])
        if kind == 'play':
            cards = np.asarray(R.counts(R.from_cards(event['cards'])))
            played[actor] += cards
            row[:15] = cards / 4
            pattern = event['pattern']
            move_kind = pattern['type']
            row[-4:-1] = ((pattern['rank'] + 1) / 15,
                          pattern.get('len', 1) / 12, len(event['cards']) / 20)
        else:
            move_kind = kind
            if kind == 'bid':
                row[-1] = event['value'] / 3
        row[15 + KINDS.index(move_kind)] = 1
        row[15 + len(KINDS) + actor] = 1
        history.append(row)
    last = np.zeros(15) if obs.last is None else R.counts(obs.last.code)
    # Online last patterns have code=0; recover actual cards from public history.
    if obs.last is not None:
        for event in reversed(latest_log(state)):
            if event.get('kind') == 'play':
                last = R.counts(R.from_cards(event['cards']))
                break
    ranks = np.stack((R.counts(obs.hand), R.counts(obs.unseen),
                      R.counts(obs.known_landlord), last, *played), axis=1) / 4
    last_actor = [float(obs.last is not None and obs.last_seat == seat) for seat in seats]
    scalars = [obs.sizes[seat] / 20 for seat in seats]
    scalars += [obs.plays[seat] / 20 for seat in seats]
    scalars += [obs.bombs / 14, obs.passes / 2, float(obs.last is None)]
    scalars += last_actor + [float(state.get('bid', 1)) / 3,
                              float(role != 0 and obs.last_seat == obs.teammate), 1.]
    rows = np.zeros((HISTORY, HISTORY_DIM), dtype=np.float32)
    mask = np.zeros(HISTORY, dtype=np.bool_)
    recent = history[-HISTORY:]
    if recent:
        rows[:len(recent)] = recent
        mask[:len(recent)] = True
    return {'ranks': torch.tensor(ranks, dtype=torch.float32),
            'scalars': torch.tensor(scalars, dtype=torch.float32),
            'history': torch.from_numpy(rows), 'mask': torch.from_numpy(mask),
            'role': role, 'actions': action_features(obs.hand, moves)}


def legal_bids(state):
    current = state.get('bid', 0) or 0
    if type(current) is not int or current not in range(4):
        raise ValueError('invalid current bid')
    bids = list(range(current + 1, 4))
    if not state.get('must_bid') or not bids:
        bids.insert(0, 0)
    return bids


def encode_bid(state, bids):
    if state.get('seat') not in (0, 1, 2) or state.get('turn', state['seat']) != state['seat']:
        raise ValueError('invalid bidding turn')
    hand = R.from_cards(state['hand'])
    if R.size(hand) != 17:
        raise ValueError('bidding requires 17 cards')
    ranks = torch.zeros(15, 7)
    ranks[:, 0] = torch.tensor(R.counts(hand)) / 4
    ranks[:, 1] = torch.tensor(R.counts(R.DECK - hand)) / 4
    scalars = torch.zeros(SCALAR_DIM)
    scalars[:3] = 17 / 20
    scalars[12] = (state.get('bid', 0) or 0) / 3
    scalars[13] = float(bool(state.get('must_bid')))
    history = torch.zeros(HISTORY, HISTORY_DIM)
    mask = torch.zeros(HISTORY, dtype=torch.bool)
    events = [event for event in latest_log(state) if event.get('kind') == 'bid'][-HISTORY:]
    for index, event in enumerate(events):
        history[index, 15 + KINDS.index('bid')] = 1
        history[index, 15 + len(KINDS) + (event['seat'] - state['seat']) % 3] = 1
        history[index, -1] = event['value'] / 3
        mask[index] = True
    actions = torch.zeros(len(bids), ACTION_DIM)
    actions[:, 15:30] = ranks[:, 0]
    actions[:, 30 + KINDS.index('bid')] = 1
    actions[:, -1] = torch.tensor(bids) / 3
    return {'ranks': ranks, 'scalars': scalars, 'history': history,
            'mask': mask, 'role': 3, 'actions': actions}


def collate(samples, device):
    result = {key: torch.stack([sample[key] for sample in samples]).to(
        device=device, dtype=torch.bool if key == 'mask' else torch.float32)
              for key in ('ranks', 'scalars', 'history', 'mask')}
    result['role'] = torch.tensor([sample['role'] for sample in samples],
                                  device=device, dtype=torch.long)
    result['actions'] = torch.cat([sample['actions'] for sample in samples]).to(
        device=device, dtype=torch.float32)
    result['owner'] = torch.repeat_interleave(
        torch.arange(len(samples)), torch.tensor([len(sample['actions']) for sample in samples])).to(device)
    return result