import asyncio
import time
import random
import copy
import json
import sys

class Engine:
    lock = asyncio.Lock()
    def __init__(self, other=None):
        self.lock = asyncio.Lock()
        if other == None:
            self.data = {
                "p": [64] * 10,
                "n": [64] * 10,
                "b": [64] * 10,
                "r": [64] * 10,
                "q": [64] * 10,
                "k": [64] * 12 }
        elif type(other) == Engine:
            self.data = copy.deepcopy(other.data)
        else:
            self.data = copy.deepcopy(other)

    def __str__(self):
        trimap = [
            0, 1, 2, 3, 3, 2, 1, 0,
            1, 4, 5, 6, 6, 5, 4, 1,
            2, 5, 7, 8, 8, 7, 5, 2,
            3, 6, 8, 9, 9, 8, 6, 3,
            3, 6, 8, 9, 9, 8, 6, 3,
            2, 5, 7, 8, 8, 7, 5, 2,
            1, 4, 5, 6, 6, 5, 4, 1,
            0, 1, 2, 3, 3, 2, 1, 0]
        colmap = [0, 1, 2, 3, 3, 2, 1, 0]
        psqt_pawn = [0] * 64
        for row in range(6):
            rv = self.data["p"][4:][row]
            for col in range(8):
                cv = self.data["p"][:4][colmap[col]]
                psqt_pawn[row * 8 + 8 + col] = (rv + cv) // 2
        psqt_king = [0] * 64
        for row in range(8):
            rv = self.data["k"][4:][row]
            for col in range(8):
                cv = self.data["k"][:4][colmap[col]]
                psqt_king[row * 8 + col] = (rv + cv) // 2
        return f"""#include "types.hh"
namespace chess {{
    extern const int psqt_avg = {sum(self.data["p"]) // len(self.data["p"])};
    extern const int psqt[6][64] = {{
        {{{ ",".join(str(i) for i in psqt_pawn) }}},
        {{{ ",".join(str(self.data["n"][i]) for i in trimap) }}},
        {{{ ",".join(str(self.data["b"][i]) for i in trimap) }}},
        {{{ ",".join(str(self.data["r"][i]) for i in trimap) }}},
        {{{ ",".join(str(self.data["q"][i]) for i in trimap) }}},
        {{{ ",".join(str(i) for i in psqt_king) }}},
    }};
}}
"""

    def mutate(self):
        for k, p in self.data.items():
            for i, v in enumerate(p):
                if random.random() < 5 / 62:
                    p[i] = random.gauss(v, v / 16)
                    p[i] = max(32, min(512, int(round(p[i]))))

    async def move(self, pos):
        async with self.lock:
            self.proc.stdin.write(f"position startpos moves { ' '.join(pos) }\n".encode())
            self.proc.stdin.write(f"go movetime 25\n".encode())
            await self.proc.stdin.drain()
            score = 0
            while True:
                line = await self.proc.stdout.readline()
                data = line.decode().strip().split(" ")
                if data[0] == "bestmove":
                    if data[1] == "0000":
                        return score
                    pos.append(data[1])
                    return None
                elif data[0] == "info":
                    score = int(data[3])
                else:
                    print(pos)
                    sys.exit(1)

    async def open(self):
        async with Engine.lock:
            with open("src/eval.cc", "w") as f:
                f.write(str(self))
            make = await asyncio.create_subprocess_exec("./make.sh")
            await make.wait()
            self.proc = await asyncio.create_subprocess_shell("./chess",
                stdout=asyncio.subprocess.PIPE,
                stdin=asyncio.subprocess.PIPE)

    async def close(self, *args):
        self.proc.stdin.write_eof()
        await self.proc.wait()
        del self.proc

    async def play(e1, e2, pos=None):
        if pos == None:
            pos = []
        score = await e1.move(pos)
        if score != None:
            return score
        return -await Engine.play(e2, e1, pos)

async def main():
    try:
        try:
            with open("temp.json", "r") as f:
                data = json.load(f)
        except:
            data = None
        engines = [Engine(data)]
        await engines[0].open()
        while True:
            while len(engines) < 6:
                engines.append(Engine(engines[0]))
                engines[-1].mutate()
                await engines[-1].open()
            results = []
            for e1 in engines:
                games = []
                for e2 in engines:
                    if e1 == e2:
                        async def null():
                            return 0
                        games.append(null())
                    else:
                        games.append(Engine.play(e1, e2))
                results.append(asyncio.gather(*games))
            results = await asyncio.gather(*results)
            scores = [0] * len(engines)
            for i, v in enumerate(results):
                for j, w in enumerate(v):
                    if w == 0:
                        scores[i] += 1
                        scores[j] += 1
                    elif w < 0:
                        scores[j] += 2
                    elif w > 0:
                        scores[i] += 2
            best_index = 0
            for i, v in enumerate(scores):
                if v > scores[best_index]:
                    best_index = i
            print(f"win {best_index}: {scores[best_index]}")
            with open("temp.json", "w") as f:
                json.dump(engines[best_index].data, f)
            engines.insert(0, engines.pop(best_index))
            while len(engines) > 1:
                await engines.pop().close()
    finally:
        for engine in engines:
            await engine.close()

if __name__ == "__main__":
    asyncio.run(main())