import numpy as np
import random

from collections import namedtuple, deque, defaultdict

import torch
import torch.nn.functional as F
import torch.optim as optim

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class ExperienceReplay():
    def __init__(self, buffer_size, batch_size):

        # Hyperparameters
        self.batch_size = batch_size
        self.buffer_size = buffer_size

        # Init replay buffer and priorities as deque lists
        self.buffer = deque(maxlen=buffer_size)
        self.experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])

    def add(self, state, action, reward, next_state, done):
        # Add an experience to the replay buffer
        self.buffer.append(self.experience(state, action, reward, next_state, done))

    def sample(self):
        # Randomly select experiences from the buffer
        idx = np.random.choice(len(self.buffer), self.batch_size)
        experiences = [self.buffer[i] for i in idx]

        states = torch.from_numpy(np.vstack([e.state for e in experiences if e is not None])).float().to(device)
        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).long().to(device)
        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(device)
        next_states = torch.from_numpy(np.vstack([e.next_state for e in experiences if e is not None])).float().to(device)
        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(device)

        return (states, actions, rewards, next_states, dones)

    def __len__(self):
        return len(self.buffer)

class PrioritizedExperienceReplay():
    def __init__(self, buffer_size, batch_size, a):

        # Hyperparameters
        self.batch_size = batch_size
        self.buffer_size = buffer_size
        self.a = a

        # Init replay buffer and priorities as deque lists
        self.buffer = deque(maxlen=buffer_size)
        self.experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])
        self.priorities = deque(np.zeros(buffer_size, np.float32), maxlen=buffer_size)

        # Index to keep where in the lists the experiences are being saved
        self.index = 0

    def add(self, state, action, reward, next_state, done):

        priorityAddArray = list(self.priorities)

        # Adding values to priority so the list starts as non zero
        priority = max(priorityAddArray) if self.buffer else 1.0

        # Add experience and priority to lists
        self.buffer.append(self.experience(state, action, reward, next_state, done))
        self.priorities.append(priority)

        # Add 1 to index every time an experience is added up to the buffer_size
        self.index = min((self.index+1), self.buffer_size)

    def sample(self, beta):

        prioritySampleArray = np.asarray(list(self.priorities), dtype = float)

        # If buffer is full then use the full list else a slice of the list is needed
        if len(self.buffer) == self.buffer_size:
            samplePriorities = prioritySampleArray
        else:
            samplePriorities = prioritySampleArray[-self.index:]

        # Calculate Sampling Probabilities
        sampleProbs = (samplePriorities)**self.a / (samplePriorities**self.a).sum()

        # Choose experiences based on Probabilities
        idx = np.random.choice(len(self.buffer), self.batch_size, p = sampleProbs, replace = True)
        experiences = [self.buffer[i] for i in idx]

        # Calculate importance sampling weights for experiences
        weights = (len(self.buffer) * sampleProbs[idx])**(-beta)
        weights = weights / weights.max()

        # Separate out S, A, R, S' and dones from experiences and transform them to a torch tensor
        states = torch.from_numpy(np.vstack([e.state for e in experiences if e is not None])).float().to(device)
        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).long().to(device)
        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(device)
        next_states = torch.from_numpy(np.vstack([e.next_state for e in experiences if e is not None])).float().to(device)
        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(device)

        return (states, actions, rewards, next_states, dones), idx, weights

    def update_priorities(self, prior, idx):
        for i in prior:
            self.priorities.append(i)

    def __len__(self):
        return len(self.buffer)


from sqlalchemy import insert, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from fastapi import HTTPException, UploadFile

from src.auth.models import User

from src.profiles.schemas import ProfileCreate, ProfileUpdate
from src.profiles.models import Profile, Avatar

from src.profiles.services.s3 import upload_avatar_to_s3

from src.profiles.exceptions.http import ProfileAlreadyExist

class ProfileDB:

    @staticmethod
    async def add_profile(profile: ProfileCreate, session: AsyncSession, user: User):
        try:
            stmt = insert(Profile).values(
                first_name=profile.first_name,
                second_name=profile.second_name,
                username=profile.username,
                user_id=user.id
            )
            await session.execute(stmt)
            await session.commit()
        except IntegrityError:
            raise ProfileAlreadyExist
        return 'success'

    @staticmethod
    async def update_profile(profile: ProfileUpdate, session: AsyncSession, user: User):
        data = {key: value for key, value in profile.dict().items() if value}
        stmt = update(Profile).where(Profile.user_id == user.id).values(**data)
        await session.execute(stmt)
        await session.commit()
        return 'success'

    @staticmethod
    async def get_my_profile(session: AsyncSession, user: User):
        query = select(Profile).where(Profile.user_id == user.id)
        result = await session.execute(query)
        return result.scalar()

class AvatarDB:

    @staticmethod
    async def get_url_file(session: AsyncSession, profile: Profile) -> Avatar:
        query = select(Avatar).where(Avatar.profile_id == profile.id)
        result = await session.execute(query)
        return result.scalar()

    @staticmethod
    async def upload_file(file: UploadFile, url: str, session: AsyncSession, profile: Profile):
        stmt = insert(Avatar).values(url=url, profile_id=profile.id)
        await session.execute(stmt)
        await session.commit()
        upload_avatar_to_s3(file, url)
        return 'success'

    @staticmethod
    async def update_upload_file(file: UploadFile, url: str, session: AsyncSession, profile: Profile):
        stmt = update(Avatar).where(Avatar.profile_id == profile.id).values(url=url)
        await session.execute(stmt)
        await session.commit()
        upload_avatar_to_s3(file, url)
        return 'update'


from typing import Iterable, Union, List, Tuple, Match

from enum import Enum
from collections import namedtuple
import re

from .__util import *

class ScanMode(Enum):
    code = 'code', None, None
    raw = 'raw', "'", "'"
    bracket_comment = 'com', '{', '}'
    star_comment = 'com', '(*', '*)'
    line_comment = 'com', '//', ''

class Comment(str):
    __slots__ = ('args',)

    @classmethod
    def full(cls, arg, open_, close):
        return open_ + arg.strip() + close

    def __new__(cls, arg: str, open_, close):
        full = cls.full(arg, open_, close)
        ret = super().__new__(cls, full)
        ret.args = arg, open_, close
        return ret

    def __repr__(self):
        return type(self).__name__ + '(' + ', '.join(self.args) + ')'

class LineParts(List[str]):
    __slots__ = ('remove_inline_comments',)

    def __new__(cls, *args, remove_inline_comments=False, **kwargs):
        return super().__new__(cls, *args, **kwargs)

    def __init__(self, *args, remove_inline_comments=False, **kwargs):
        super().__init__(*args, **kwargs)
        self.remove_inline_comments = remove_inline_comments

    @property
    def last_is_comment(self):
        for i in reversed(self):
            if isinstance(i, Comment):
                return True
            if isblank(i):
                continue
            return False
        return False

    def append(self, o, mode=None):
        if mode and mode.value[0] == 'com':
            o = Comment(o, *mode.value[1:])
        if not isinstance(o, Comment) and not o.isspace():
            while self.last_is_comment:
                if self.remove_inline_comments:
                    self.pop()
                    continue
                raise NotImplementedError('inline comments must be at the end of a line')
        super().append(o)

    def join(self):
        return ''.join(str(i) for i in self)

Region = namedtuple('Region', 'start_pattern end_pattern mode')
regions = [
    Region(re.compile("'"), re.compile(r"(?<!\\)'"), ScanMode.raw),
    Region(re.compile('{(?!\$)'), re.compile('}'), ScanMode.bracket_comment),
    Region(re.compile('\(\*'), re.compile('\*\)'), ScanMode.star_comment),
    Region(re.compile('//'), re.compile('$'), ScanMode.line_comment),
]

def next_region_ind(line):
    """
    returns the index of the first matching region. If not found, returns the line's length
    >>> next_region_ind("abc{d'ef")
    3
    >>> next_region_ind("re'ge{x")
    2
    >>> next_region_ind("regex")
    5
    """
    matches = [reg.start_pattern.search(line) for reg in regions]
    match = min(matches,
                key=lambda m: float('inf') if m is None else m.start())
    if match is None:
        return len(line)

    return match.start()

def match_region(line) -> Tuple[Region, Match]:
    """
    returns the region that matches the first characters in the line, raises an error if there is more or less than one
    matching region. Also returns the match
    """
    matches = [(reg, reg.start_pattern.match(line)) for reg in regions]
    matches = [(k, v) for (k, v) in matches if v]
    if len(matches) != 1:
        raise AssertionError(f'expected 1 match, got {len(matches)}')
    return matches[0]

def filter_multiline_comments(source: Iterable[str], remove_inline_comments=False):
    """
    all this function does is wrap all lines in between multiline comments with a seperate multiline comment
        {a
        b
        c}
    to:
        {a}
        {b}
        {c}
    >>> assert list(filter_multiline_comments(["abcd","efgh"])) == ["abcd","efgh"]
    >>> assert list(filter_multiline_comments(["abc{def","def","def}"])) == ["abc{def}","{def}","{def}"]
    >>> assert list(filter_multiline_comments(["abc{def","   def","def}"])) == ["abc{def}","   {def}","{def}"]
    >>> assert list(filter_multiline_comments(["my mother said 'i need {some} dire frui{{t'" \
                                                ,"{she didn't actually","say","that}" \
                                                ,"but you {don't care} {do you?}"])) \
                                                 == ["my mother said 'i need {some} dire frui{{t'" \
                                                 ,"{she didn't actually}","{say}","{that}" \
                                                 ,"but you {don't care} {do you?}"]
    >>> assert list(filter_multiline_comments(['hi //there','how are you?','{i','am}','//fi{ne','thanks', \
                                                '{no //problem}'])) \
                                                == ['hi //there', 'how are you?', '{i}', '{am}', '//fi{ne', 'thanks',\
                                                 '{no //problem}']
    """

    region: Union[Region, type(...), None] = None
    # the region we are in, None means we are in code
    # ... means the next token is the start of a region, and we need to find out which
    for line in source:
        line_parts = LineParts(remove_inline_comments=remove_inline_comments)
        if region is not None:
            indent = re.match('\s*', line)
            line_parts.append(indent.group())
            line = line[indent.end():]
        capture = ''
        while line:
            if region is None:
                code_scan_length = next_region_ind(line)
                if code_scan_length == 0:
                    region = ...
                else:
                    line_parts.append(line[:code_scan_length])
                    line = line[code_scan_length:]
            elif region is ...:
                region, match = match_region(line)
                if region.mode.value[0] == 'com':
                    capture = ''
                elif region.mode == ScanMode.raw:
                    capture = match[0]
                else:
                    raise AssertionError('unhandled region mode: ' + str(region.mode))
                line = line[match.end():]
            else:
                # regular region (string or comment)
                end_token_match = region.end_pattern.search(line)
                mode = region.mode
                if not end_token_match:
                    # the rest of the line is the current region (it continues unto the next line but we split it here)
                    # technically, this is only allowed if we're in a comment, but I'm not about to start throwing
                    # pascal errors
                    capture += line
                    line = ''
                else:
                    if region.mode.value[0] == 'com':
                        end_cap_ind = end_token_match.start()
                    elif region.mode == ScanMode.raw:
                        end_cap_ind = end_token_match.end()
                    else:
                        raise AssertionError('unhandled region mode: ' + str(region.mode))
                    capture += line[:end_cap_ind]
                    line = line[end_token_match.end():]
                    region = None
                # either way, this region is done
                line_parts.append(capture, mode)
                capture = ''
        yield line_parts.join()


from typing import (
    List,
    Tuple,
    Type,
    Union,
)

import logging
from decimal import Decimal

from sqlalchemy import select
from sqlalchemy.sql import func

import xquery.cache
import xquery.db
import xquery.db.orm as orm
import xquery.event.indexer
from xquery.contract import (
    png_factory,
    psys_factory,
)
from .processor import EventProcessorStage

log = logging.getLogger(__name__)

class EventProcessorStageExchange_Count(EventProcessorStage):

    def __init__(self, db: xquery.db.FusionSQL, cache: xquery.cache.Cache, factory_address: str) -> None:
        """
        Exchange processor stage to aggregate transaction and pair counts.

        Computes the following data:
          - factory.pairCount
          - factory.txCount
          - pair.txCount
          - pair.totalSupply
          - pair.volumeToken0
          - pair.volumeToken1
          - token.txCount
          - token.tradeVolume

        Depends on:
          - None

        Note: Should be run in a single job

        :param db: database service
        :param cache: cache service
        """
        super().__init__(db, cache)

        self.factory_address = factory_address

    def _aggregate_factory(self, session, start_block: int, end_block: int) -> List[orm.Base]:
        factory = session.execute(
            select(orm.Factory)
                .filter(orm.Factory.address == self.factory_address)
                .limit(1)
        ).scalar()
        assert factory is not None

        # update the pair count
        factory.pairCount += (
            session.query(func.count(orm.Pair.id))
                .join(orm.Block)
                .filter(orm.Block.number.between(start_block, end_block))
        ).scalar()

        # update the tx count (actually more of an event counter)
        # sum (mint, burn, swap) between (start_block, end_block)
        factory.txCount += (
            session.query(func.count(orm.Mint.id))
                .join(orm.Transaction)
                .join(orm.Block)
                .filter(orm.Block.number.between(start_block, end_block))
        ).scalar()

        factory.txCount += (
            session.query(func.count(orm.Burn.id))
                .join(orm.Transaction)
                .join(orm.Block)
                .filter(orm.Block.number.between(start_block, end_block))
        ).scalar()

        factory.txCount += (
            session.query(func.count(orm.Swap.id))
                .join(orm.Transaction)
                .join(orm.Block)
                .filter(orm.Block.number.between(start_block, end_block))
        ).scalar()

        return [factory]

    def _aggregate_pair(self, session, start_block: int, end_block: int) -> List[orm.Base]:
        pairs = session.execute(
            select(orm.Pair)
                .order_by(orm.Pair.id)
        ).yield_per(1000).scalars()

        # update the tx count (actually more of an event counter)
        # sum (mint, burn, swap) between (start_block, end_block)
        objects = []
        for pair in pairs:
            mint_count, mint_value, mint_fee_value = (
                session.query(func.count(orm.Mint.id), func.sum(orm.Mint.liquidity), func.sum(orm.Mint.feeLiquidity))
                    .filter(orm.Mint.pair_address == pair.address)
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).one()
            pair.txCount += mint_count
            pair.totalSupply += Decimal(0) if mint_value is None else mint_value
            pair.totalSupply += Decimal(0) if mint_fee_value is None else mint_fee_value

            burn_count, burn_value, burn_fee_value = (
                session.query(func.count(orm.Burn.id), func.sum(orm.Burn.liquidity), func.sum(orm.Burn.feeLiquidity))
                    .filter(orm.Burn.pair_address == pair.address)
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).one()
            pair.txCount += burn_count
            pair.totalSupply -= Decimal(0) if burn_value is None else burn_value
            pair.totalSupply += Decimal(0) if burn_fee_value is None else burn_fee_value
            assert pair.totalSupply >= Decimal(0)

            swap_count, swap_value0, swap_value1 = (
                session.query(func.count(orm.Swap.id), func.sum(orm.Swap.amount0Out + orm.Swap.amount0In), func.sum(orm.Swap.amount1Out + orm.Swap.amount1In))
                    .filter(orm.Swap.pair_address == pair.address)
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).one()
            pair.txCount += swap_count
            pair.volumeToken0 += Decimal(0) if swap_value0 is None else swap_value0
            pair.volumeToken1 += Decimal(0) if swap_value1 is None else swap_value1

            objects.append(pair)

        return objects

    def _aggregate_token(self, session, start_block: int, end_block: int) -> List[orm.Base]:
        tokens = session.execute(
            select(orm.Token)
                .order_by(orm.Token.id)
        ).yield_per(1000).scalars()

        # update the tx count (actually more of an event counter)
        # sum (mint, burn, swap) between (start_block, end_block)
        objects = []
        for token in tokens:
            token.txCount += (
                session.query(func.count(orm.Mint.id))
                    .join(orm.Pair)
                    .filter((orm.Pair.token0_address == token.address) | (orm.Pair.token1_address == token.address))
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).scalar()

            token.txCount += (
                session.query(func.count(orm.Burn.id))
                    .join(orm.Pair)
                    .filter((orm.Pair.token0_address == token.address) | (orm.Pair.token1_address == token.address))
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).scalar()

            swap_count, swap_value = (
                session.query(func.count(orm.Swap.id), func.sum(orm.Swap.amount0Out + orm.Swap.amount0In))
                    .join(orm.Pair)
                    .filter(orm.Pair.token0_address == token.address)
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).one()
            token.txCount += swap_count
            token.tradeVolume += Decimal(0) if swap_value is None else swap_value

            swap_count, swap_value = (
                session.query(func.count(orm.Swap.id), func.sum(orm.Swap.amount1Out + orm.Swap.amount1In))
                    .join(orm.Pair)
                    .filter(orm.Pair.token1_address == token.address)
                    .join(orm.Transaction)
                    .join(orm.Block)
                    .filter(orm.Block.number.between(start_block, end_block))
            ).one()
            token.txCount += swap_count
            token.tradeVolume += Decimal(0) if swap_value is None else swap_value

            objects.append(token)

        return objects

    @classmethod
    def setup(cls, db: xquery.db.FusionSQL, first_block: int) -> Union[List[orm.Base], List[Tuple[Type[orm.Base], List[dict]]]]:
        return []

    def process(self, start_block: int, end_block: int) -> Union[List[orm.Base], List[Tuple[Type[orm.Base], List[dict]]]]:
        objects = []
        with self._db.session() as session:
            objects.extend(self._aggregate_factory(session, start_block, end_block))
            objects.extend(self._aggregate_pair(session, start_block, end_block))
            objects.extend(self._aggregate_token(session, start_block, end_block))

        return objects

class EventProcessorStageExchangePangolin_Count(EventProcessorStageExchange_Count):

    def __init__(self, db: xquery.db.FusionSQL, cache: xquery.cache.Cache):
        """
        Event processor stage "count" for the Pangolin Exchange (on AVAX)
        """
        super().__init__(
            db=db,
            cache=cache,
            factory_address=png_factory.address,
        )

class EventProcessorStageExchangePegasys_Count(EventProcessorStageExchange_Count):

    def __init__(self, db: xquery.db.FusionSQL, cache: xquery.cache.Cache):
        """
        Event processor stage "count" for the Pegasys Exchange (on SYS)
        """
        super().__init__(
            db=db,
            cache=cache,
            factory_address=psys_factory.address,
        )