import numpy as np
import pandas as pd
import aljpy.download
import sqlalchemy
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, Float, String, ForeignKey, create_engine
from pavlov import runs, storage
import ast
from tqdm.auto import tqdm
from pathlib import Path
from contextlib import contextmanager
from logging import getLogger

log = getLogger(__name__)

# First modern run
FIRST_RUN = pd.Timestamp('2021-02-03 12:47:26.557749+00:00')

DATABASE = Path('output/experiments/eval/database.sql')

URL = 'https://f002.backblazeb2.com/file/boardlaw/output/experiments/eval/database.sql'

Base = declarative_base()
class Run(Base):
    """One row per training run. A training run corresponds to one folder in `output/pavlov`. They
    are usually generated by `main.run`."""

    __tablename__ = 'runs'

    run = Column(String, primary_key=True)
    """Name of the training run, typically in the format `date - nickname`. This corresponds to 
    the folder name that the run is stored in by `pavlov`."""

    description = Column(String)
    """Description of the group the training run is in, typically in the format `groupname/boardsize`.
    All the data worth looking at has the description `bee/`; the other descriptions are for abandonned
    experimental sets.
    """

    boardsize = Column(Integer)
    """Size of the board the run was trained on."""

    width = Column(Integer)
    """Width of the network that was trained."""

    depth = Column(Integer)
    """Depth of the network that was trained."""

    nodes = Column(Integer)
    """Number of nodes in the MCTS during training."""


class Snap(Base):
    """One row per snapshot. A snapshot is a version of the model from a specific point in a training run. They 
    are usually generated by `main` and `storage.LogarithmicStorer`."""
    __tablename__ = 'snaps'

    id = Column(Integer, primary_key=True)
    """Auto-incremented snapshot ID."""

    run = Column(String, ForeignKey('runs.run'))
    """The training run the snapshot was taken during."""

    idx = Column(Integer)
    """The index of the snapshot, ranging from 0 for the first snapshot in a run to (typically) 20 for the last one."""

    samples = Column(Float)
    """The number of training samples seen by the time the snapshot was taken. See the `main` module and 
    `storage.LogarithmicStorer.step` method for details. """ 

    flops = Column(Float)
    """The number of floating point operations done by the time the snapshot was taken. See the `main` module 
    and `storage.LogarithmicStorer.step` method for details. """


class Agent(Base):
    """One row per agent. An agent is a snapshot plus a few parameters describing the test-time search."""
    __tablename__ = 'agents'

    id = Column(Integer, primary_key=True)
    """Auto-incremented agent ID."""

    snap = Column(Integer, ForeignKey('snaps.id'))
    """ID of the snapshot this agent is based on.""" 

    nodes = Column(Integer)
    """Number of nodes in the test-time search tree."""

    c = Column(Float)
    """The policy-vs-value test-time search coefficient, c_puct."""


class Trial(Base):
    """One row per trial. A trial is the outcome of a set of games between two agents. 
    These are usually generated by `arena.best` and `arena.neural`. Be aware that a pair of agents 
    usually have more than one row here, since trials were conducted piecemeal! In particular, 
    every pair of agents has at least one row where one agent plays black, and another row where they're 
    swapped and that agent instead plays white."""
    __tablename__ = 'trials'

    id = Column(Integer, primary_key=True)
    """Auto-incremented trial ID."""

    black_agent = Column(Integer, ForeignKey('agents.id'))
    """ID of the agent playing black."""

    white_agent = Column(Integer, ForeignKey('agents.id'))
    """ID of the agent playing black."""

    black_wins = Column(Integer)
    """Number of wins by black. Sum of this and white wins gives you the number of games."""

    white_wins = Column(Integer)
    """Number of wins by white. Sum of this and black wins gives you the number of games."""

    moves = Column(Integer)
    """Number of moves in the games, total"""

    times = Column(Integer)
    """Amount of time used, total. See `arena.common.gather` for details on how this is calculated."""


class MohexTrial(Base):
    """One row per mohex trial. A mohex trial is the outcome of a set of games between an agent and MoHex, and 
    is used to calibrate the top-ranking agents. These are usually generated by `arena.mohex`.""" 
    __tablename__ = 'mohex_trials'

    id = Column(Integer, primary_key=True)
    """Auto-incremented trial ID."""

    black_agent = Column(Integer, ForeignKey('agents.id'), nullable=True)
    """ID of the agent playing black. Null if the agent was MoHex."""

    white_agent = Column(Integer, ForeignKey('agents.id'), nullable=True)
    """ID of the agent playing white. Null if the agent was MoHex."""

    black_wins = Column(Integer)
    """Number of wins by black. Sum of this and white wins gives you the number of games."""

    white_wins = Column(Integer)
    """Number of wins by white. Sum of this and black wins gives you the number of games."""

    moves = Column(Integer)
    """Number of moves in the games, total"""

    times = Column(Integer)
    """Amount of time used, total. See `arena.common.gather` for details on how this is calculated."""


class NoiseScale(Base):
    __tablename__ = 'noise_scales'

    id = Column(Integer, primary_key=True)
    agent_id = Column(Integer, ForeignKey('agents.id'))
    kind = Column(String)
    mean_sq = Column(Float)
    sq_mean = Column(Float)
    variance = Column(Float)
    n_params = Column(Float)
    batch_size = Column(Float)
    batches = Column(Float) 

_engine = None
@contextmanager
def connection():
    if not DATABASE.parent.exists():
        DATABASE.parent.mkdir(exist_ok=True, parents=True)
    if not DATABASE.exists():
        log.info('Downloading the SQL database')
        DATABASE.write_bytes(aljpy.download.download(URL))

    global _engine
    _engine = create_engine('sqlite:///' + str(DATABASE)) if _engine is None else _engine
    with _engine.connect() as conn:
        yield conn

def create():
    with connection() as conn:
        Base.metadata.create_all(conn.engine)

        conn.execute('''
            create view agents_details as
            select 
                agents.id, agents.nodes as test_nodes, agents.c as test_c,
                snaps.id as snap_id, snaps.samples, snaps.flops as train_flops, snaps.idx, 
                runs.run, runs.description, runs.boardsize, runs.width, runs.depth, runs.nodes as train_nodes
            from agents
                inner join snaps on (agents.snap == snaps.id)
                inner join runs on (snaps.run == runs.run)''')

def run_data():
    r = runs.pandas().loc[lambda df: df._created >= FIRST_RUN]
    params = r.params.dropna().apply(pd.Series).reindex(r.index)
    insert = pd.concat([r.index.to_series().to_frame('run'), r[['description']], params[['boardsize', 'width', 'depth', 'nodes']]], 1)
    insert['nodes'] = insert.nodes.fillna(64)
    return insert.reset_index(drop=True)

def snapshot_data(new_runs):
    snapshots = {}
    for _, r in tqdm(list(new_runs.iterrows()), desc='snapshots'):
        for i, s in storage.snapshots(r.run).items():
            stored = storage.load_snapshot(r.run, i)
            if 'n_samples' in stored:
                snapshots[r.run, i] = {
                    'samples': stored['n_samples'], 
                    'flops': stored['n_flops']}
    snapshots = (pd.DataFrame.from_dict(snapshots, orient='index')
                    .rename_axis(index=('run', 'idx'))
                    .reset_index())
    # snapshots['id'] = snapshots.index.to_series()
    return snapshots

def refresh_runs():
    current = run_data()
    with connection() as conn:
        old = pd.read_sql_query('select * from runs', conn)
        new = current[~current.run.isin(old.run)]
        new.to_sql('runs', conn, index=False, if_exists='append')

        snaps = snapshot_data(new)
        snaps.to_sql('snaps', conn, index=False, if_exists='append')

    return new

def create_agents(rs, test_nodes=64, c=1/16, dry_run=False):
    with connection() as conn:
        snaps = pd.read_sql_query('select * from snaps', conn, index_col='id')

        agents = (snaps.index
            [snaps.run.isin(rs.run)]
            .to_frame(name='snap')
            .reset_index(drop=True))
        agents['nodes'] = test_nodes
        agents['c'] = c

        old_agents = pd.read_sql_query('select * from agents', conn)
        old_agents = pd.merge(
            agents.rename_axis(index='idx').reset_index(), 
            old_agents, 
            on=['snap', 'nodes', 'c'], how='inner')
        new_agents = agents[~agents.index.isin(old_agents.idx)]

        if dry_run:
            print(new_agents)
        else:
            new_agents.to_sql('agents', conn, index=False, if_exists='append')

def execute(sql, *args, **kwargs):
    with connection() as conn:
        return conn.execute(sql, *args, **kwargs)

def query(sql, *args, **kwargs):
    if len(args) > 0 and 'params' not in kwargs:
        kwargs['params'] = args
    with connection() as conn:
        return pd.read_sql_query(sql, conn, **kwargs)

def query_one(sql, *args, **kwargs):
    df = query(sql, *args, **kwargs)
    if len(df) == 1:
        return df.iloc[0]
    raise ValueError(f'query_one returned {len(df)} results')

def agent_query():
    return query('''select * from agents_details''', index_col='id')

def trial_query(boardsize=None, desc='%', test_nodes=None):
    q = '''
        select trials.* 
        from trials 
            inner join agents_details as black
                on (trials.black_agent == black.id)
            inner join agents_details as white
                on (trials.white_agent == white.id)
        where 
            (black.description like ?) and (white.description like ?)
            '''
    params = (desc, desc)
    
    if test_nodes is not None:
        q += '\nand (black.test_nodes == ?) and (white.test_nodes == ?)'
        params += (test_nodes, test_nodes)
    if boardsize is not None:
        q += '\nand (black.boardsize == ?) and (white.boardsize == ?)'
        params += (int(boardsize), int(boardsize)) 
    
    return query(q, index_col='id', params=params)

def save_trials(results):
    rows = []
    for r in results:
        rows.append({
            'black_agent': r.names[0],
            'white_agent': r.names[1],
            'black_wins': r.wins[0],
            'white_wins': r.wins[1],
            'moves': r.moves,
            'times': r.times})
    rows = pd.DataFrame(rows)
    with connection() as conn:
        rows.to_sql('trials', conn, index=False, if_exists='append')

def save_mohex_trials(results):
    rows = []
    for r in results:
        assert sum(n is None for n in r.names) == 1, 'One agent should be MoHex'
        rows.append({
            'black_agent': r.names[0],
            'white_agent': r.names[1],
            'black_wins': r.wins[0],
            'white_wins': r.wins[1],
            'moves': r.moves,
            'times': r.times})
    rows = pd.DataFrame(rows)
    with connection() as conn:
        rows.to_sql('mohex_trials', conn, index=False, if_exists='append')

def mohex_trial_query(boardsize, desc='%'):
    return query('''
        select mohex_trials.* 
        from mohex_trials 
            left join agents_details as black
                on (mohex_trials.black_agent == black.id)
            left join agents_details as white
                on (mohex_trials.white_agent == white.id)
        where 
            ((black.boardsize == ?) or (white.boardsize == ?)) and 
            ((black.description like ?) or (white.description like ?))''', index_col='id', params=(int(boardsize), int(boardsize), desc, desc))

def save_noise_scale(rows):
    if isinstance(rows, pd.Series):
        rows = pd.DataFrame([rows])
    with connection() as conn:
        rows.to_sql('noise_scales', conn, index=False, if_exists='append')

def file_change_counter():
    # https://www.sqlite.org/fileformat.html
    counter = DATABASE.open('rb').read(30)[24:28]
    dt = np.dtype('int32').newbyteorder('>')
    return int(np.frombuffer(counter, dtype=dt))