use crate::bitboard::*;
use crate::board::*;
use crate::rays::Rays;

const BISHOP_SHR: u8 = 55;
const ROOK_SHR: u8 = 52;

pub(crate) struct Magics {
    bishop: BySquare<Magic>,
    rook: BySquare<Magic>,
    table: Box<[Bitboard]>,
}

#[derive(Clone, Copy)]
struct Magic {
    premask: Bitboard,
    factor: u64,
    offset: isize,
}

impl Magics {
    pub(crate) fn compute(rays: &Rays) -> Self {
        let mut data = Vec::new();

        let mut aux =
            |shr,
             factors: fn(Square) -> u64,
             make_table: fn(&Rays, Square) -> (Bitboard, Vec<(Bitboard, Bitboard)>)| {
                BySquare::new(|square| {
                    let (premask, table) = make_table(rays, square);
                    let factor = factors(square);
                    let offset = fill_table(&mut data, shr, factor, table);
                    Magic {
                        premask,
                        factor,
                        offset,
                    }
                })
            };

        let bishop = aux(BISHOP_SHR, bishop_factors, make_bishop_table);
        let rook = aux(ROOK_SHR, rook_factors, make_rook_table);

        let mut table = Box::new_uninit_slice(data.len());
        for (i, entry) in data.into_iter().enumerate() {
            table[i].write(entry);
        }

        Self {
            bishop,
            rook,
            table: unsafe { table.assume_init() },
        }
    }

    #[inline]
    pub(crate) fn bishop(&self, square: Square, blockers: Bitboard) -> Bitboard {
        unsafe { self.get_unchecked(BISHOP_SHR, *self.bishop.get(square), blockers) }
    }

    #[inline]
    pub(crate) fn rook(&self, square: Square, blockers: Bitboard) -> Bitboard {
        unsafe { self.get_unchecked(ROOK_SHR, *self.rook.get(square), blockers) }
    }

    #[inline]
    unsafe fn get_unchecked(&self, shr: u8, magic: Magic, blockers: Bitboard) -> Bitboard {
        let Magic {
            premask,
            factor,
            offset,
        } = magic;
        *self.table.get_unchecked(
            ((hash(shr, factor, blockers | premask) as isize).unchecked_add(offset)) as usize,
        )
    }
}

fn fill_table(
    data: &mut Vec<Bitboard>,
    shr: u8,
    factor: u64,
    table: Vec<(Bitboard, Bitboard)>,
) -> isize {
    let offset = data.len() as isize
        - table
            .iter()
            .map(|(x, _)| hash(shr, factor, *x) as isize)
            .min()
            .unwrap();
    for (x, y) in &table {
        let i = (hash(shr, factor, *x) as isize + offset) as usize;
        while data.len() <= i {
            data.push(Bitboard::new());
        }
        if data[i] != Bitboard::new() && data[i] != *y {
            panic!();
        }
        data[i] = *y;
    }
    offset
}

fn make_bishop_table(rays: &Rays, square: Square) -> (Bitboard, Vec<(Bitboard, Bitboard)>) {
    let mut premask = Bitboard::new();
    for direction in [
        Direction::NorthWest,
        Direction::SouthWest,
        Direction::SouthEast,
        Direction::NorthEast,
    ] {
        premask |= rays.ray(square, direction);
    }
    premask &= !Rank::First.bitboard();
    premask &= !Rank::Eighth.bitboard();
    premask &= !File::A.bitboard();
    premask &= !File::H.bitboard();

    let mut table = make_table(premask, |blockers| {
        let mut res = Bitboard::new();
        for direction in [
            Direction::NorthWest,
            Direction::SouthWest,
            Direction::SouthEast,
            Direction::NorthEast,
        ] {
            res |= rays.blocked(square, direction, blockers);
        }
        res
    });

    premask = !premask;
    for (x, _) in &mut table {
        *x |= premask;
    }

    (premask, table)
}

fn make_rook_table(rays: &Rays, square: Square) -> (Bitboard, Vec<(Bitboard, Bitboard)>) {
    let mut premask = Bitboard::new();
    premask |= rays.ray(square, Direction::North) & !Rank::Eighth.bitboard();
    premask |= rays.ray(square, Direction::West) & !File::A.bitboard();
    premask |= rays.ray(square, Direction::South) & !Rank::First.bitboard();
    premask |= rays.ray(square, Direction::East) & !File::H.bitboard();

    let mut table = make_table(premask, |blockers| {
        let mut res = Bitboard::new();
        for direction in [
            Direction::North,
            Direction::West,
            Direction::South,
            Direction::East,
        ] {
            res |= rays.blocked(square, direction, blockers);
        }
        res
    });

    premask = !premask;
    for (x, _) in &mut table {
        *x |= premask;
    }

    (premask, table)
}

fn make_table<T, F>(premask: Bitboard, f: F) -> Vec<(Bitboard, T)>
where
    F: Fn(Bitboard) -> T,
{
    let mut res = Vec::new();
    let mut subset: u64 = 0;
    loop {
        subset = subset.wrapping_sub(premask.0) & premask.0;
        let x = Bitboard(subset);
        let y = f(x);
        res.push((x, y));
        if subset == 0 {
            break;
        }
    }
    res
}

fn hash(shr: u8, factor: u64, x: Bitboard) -> usize {
    (x.0.wrapping_mul(factor) >> shr) as usize
}

fn bishop_factors(square: Square) -> u64 {
    match square {
        Square::A1 => 0x0000404040404040,
        Square::B1 => 0x0040C100081000E8,
        Square::C1 => 0x0000401020200000,
        Square::D1 => 0x0040802004000000,
        Square::E1 => 0x10403C0180000000,
        Square::F1 => 0x0040210100800000,
        Square::G1 => 0x0068104002008000,
        Square::H1 => 0x0048082080040080,
        Square::A2 => 0x0000004040404040,
        Square::B2 => 0x0000002020202020,
        Square::C2 => 0x00040080184001E4,
        Square::D2 => 0x0040008020040000,
        Square::E2 => 0x1040003C01800000,
        Square::F2 => 0x0078002001008000,
        Square::G2 => 0x0068001040020080,
        Square::H2 => 0x0068000820010040,
        Square::A3 => 0x0000400080808080,
        Square::B3 => 0x0000200040404040,
        Square::C3 => 0x0000400080808080,
        Square::D3 => 0x0000200200801000,
        Square::E3 => 0x0060200100080000,
        Square::F3 => 0x0000100021C60021,
        Square::G3 => 0x0000040010410040,
        Square::H3 => 0x0000020008208020,
        Square::A4 => 0x0000804000810100,
        Square::B4 => 0x0000402000408080,
        Square::C4 => 0x0000040800802080,
        Square::D4 => 0x000020100C010020,
        Square::E4 => 0x0000840000802000,
        Square::F4 => 0x0001801800240010,
        Square::G4 => 0x0000080800104100,
        Square::H4 => 0x0000040400082080,
        Square::A5 => 0x0000010278010040,
        Square::B5 => 0x000000813C004040,
        Square::C5 => 0x000001027A010040,
        Square::D5 => 0x0000018180280200,
        Square::E5 => 0x0000204018003080,
        Square::F5 => 0x0000202040008040,
        Square::G5 => 0x0000101010002080,
        Square::H5 => 0x0000080808001040,
        Square::A6 => 0x0000004100F90080,
        Square::B6 => 0x0000002080BC0040,
        Square::C6 => 0x0000004103440080,
        Square::D6 => 0x0000000080FD0080,
        Square::E6 => 0x0000020040100100,
        Square::F6 => 0x0000404040400080,
        Square::G6 => 0x000000206027D010,
        Square::H6 => 0x00000008400DE806,
        Square::A7 => 0x0000002101007200,
        Square::B7 => 0x0000001041003900,
        Square::C7 => 0x080000000F8080A0,
        Square::D7 => 0x0000000008003FC0,
        Square::E7 => 0x0000000100202000,
        Square::F7 => 0x0000004040802000,
        Square::G7 => 0x00000060401043D0,
        Square::H7 => 0x00000020200413F0,
        Square::A8 => 0x1400000F00410088,
        Square::B8 => 0x0000000010410039,
        Square::C8 => 0x000080000800807E,
        Square::D8 => 0x000C69003008003F,
        Square::E8 => 0x0000000001002020,
        Square::F8 => 0x0000000040408020,
        Square::G8 => 0x001980004010801F,
        Square::H8 => 0x0000404040404040,
    }
}

fn rook_factors(square: Square) -> u64 {
    match square {
        Square::A1 => 0x002000A28110000C,
        Square::B1 => 0x0018000C01060001,
        Square::C1 => 0x0040080010004004,
        Square::D1 => 0x0028004084200028,
        Square::E1 => 0x0030018000900300,
        Square::F1 => 0x0020008020010202,
        Square::G1 => 0x001800410080001F,
        Square::H1 => 0x0068006801040004,
        Square::A2 => 0x000028010114000A,
        Square::B2 => 0x00000C0083000600,
        Square::C2 => 0x0000080401020008,
        Square::D2 => 0x0000200200040020,
        Square::E2 => 0x0000200100020020,
        Square::F2 => 0x00001800C0006018,
        Square::G2 => 0x0000180070400018,
        Square::H2 => 0x0000180030640018,
        Square::A3 => 0x00300018010C0004,
        Square::B3 => 0x0004001000080010,
        Square::C3 => 0x0001000804020008,
        Square::D3 => 0x0002002004002002,
        Square::E3 => 0x0001002002002001,
        Square::F3 => 0x0001001000801040,
        Square::G3 => 0x0000004040008001,
        Square::H3 => 0x0000802000200040,
        Square::A4 => 0x0040200010080008,
        Square::B4 => 0x0000080010040010,
        Square::C4 => 0x0001020008040008,
        Square::D4 => 0x0000020020040020,
        Square::E4 => 0x0000010020020020,
        Square::F4 => 0x0000008020010020,
        Square::G4 => 0x0000008020200040,
        Square::H4 => 0x0000200020004081,
        Square::A5 => 0x0000081C00380020,
        Square::B5 => 0x0000080400100010,
        Square::C5 => 0x0000400880410010,
        Square::D5 => 0x0000200200200400,
        Square::E5 => 0x0000200100200200,
        Square::F5 => 0x0000200080200100,
        Square::G5 => 0x0000008000404001,
        Square::H5 => 0x0000802000200040,
        Square::A6 => 0x0000010B14002800,
        Square::B6 => 0x0000030086000C00,
        Square::C6 => 0x0000084040804200,
        Square::D6 => 0x0000020004002020,
        Square::E6 => 0x0000009001803003,
        Square::F6 => 0x0000004001004002,
        Square::G6 => 0x000000100800A804,
        Square::H6 => 0x000000082800D002,
        Square::A7 => 0x00000109040200A8,
        Square::B7 => 0x000000808A050014,
        Square::C7 => 0x0000004048038018,
        Square::D7 => 0x0000020020040020,
        Square::E7 => 0x0000002030018030,
        Square::F7 => 0x0000001800E08018,
        Square::G7 => 0x0000000810580050,
        Square::H7 => 0x0000000C04600050,
        Square::A8 => 0x0000001020891046,
        Square::B8 => 0x00000080090015C1,
        Square::C8 => 0x0000004005489101,
        Square::D8 => 0x0000040810204002,
        Square::E8 => 0x000C040810002022,
        Square::F8 => 0x0008000404883002,
        Square::G8 => 0x0008000400548802,
        Square::H8 => 0x0000000224104486,
    }
}