from django.conf.urls import patterns, url

from keymgmt.views import (
    HomeView,
    EnvironmentList,
    EnvironmentCreate,
    EnvironmentDelete,
    EnvironmentDetail,
    EnvironmentUpdate,
    SSHKeyList,
    SSHKeyCreate,
    SSHKeyUpdate,
    SSHKeyDelete,
    SSHKeyDetail,
    GroupList,
    GroupCreate,
    GroupDelete,
    GroupDetail,
    HostList,
    HostCreate,
    HostDetail,
    HostDelete,
    HostUpdate,
    SSHKeyringList,
    SSHKeyringCreate,
    SSHKeyringUpdate,
    SSHKeyringDelete,
    SSHKeyringDetail,
    SSHAccountList,
    SSHAccountDetail,
    SSHAccountCreateEnvironment,
    SSHAccountCreateHost,
    SSHAccountCreateGroup,
    SSHAccountUpdate,
    SSHAccountDelete,
    SSHAccountKeyUpdate,
    SSHAccountAvailableList,
    SSHAccountAvailableDelete,
    GroupRuleCreate,
    GroupRuleDelete,
    GroupRuleUpdate,
    AuditKey2Access,
    api_get_keys,
)

urlpatterns = patterns('',
    url(r'^environment/$', EnvironmentList.as_view(), name='environment_list'),
    url(r'^environment/add/$', EnvironmentCreate.as_view(), name='environment_create'),
    url(r'^environment/(?P<pk>[0-9]+)/delete/$', EnvironmentDelete.as_view(), name='environment_delete'),
    url(r'^environment/(?P<pk>[0-9]+)/update/$', EnvironmentUpdate.as_view(), name='environment_update'),
    url(r'^environment/(?P<pk>[0-9]+)/$', EnvironmentDetail.as_view(), name='environment_detail'),

    url(r'^sshkey/$', SSHKeyList.as_view(), name='sshkey_list'),
    url(r'^sshkey/add/$', SSHKeyCreate.as_view(), name='sshkey_create'),
    url(r'^sshkey/(?P<pk>[0-9]+)/delete/$', SSHKeyDelete.as_view(), name='sshkey_delete'),
    url(r'^sshkey/(?P<pk>[0-9]+)/update/$', SSHKeyUpdate.as_view(), name='sshkey_update'),
    url(r'^sshkey/(?P<pk>[0-9]+)/$', SSHKeyDetail.as_view(), name='sshkey_detail'),

    url(r'^group/$', GroupList.as_view(), name='group_list'),
    url(r'^group/add/$', GroupCreate.as_view(), name='group_create'),
    url(r'^group/(?P<pk>[0-9]+)/delete/$', GroupDelete.as_view(), name='group_delete'),
    url(r'^group/(?P<pk>[0-9]+)/$', GroupDetail.as_view(), name='group_detail'),

    url(r'^group/(?P<pk>[0-9]+)/rule/add/$', GroupRuleCreate.as_view(), name='group_rule_create'),
    url(r'^group/(?P<group_pk>[0-9]+)/rule/(?P<pk>[0-9]+)/delete/$', GroupRuleDelete.as_view(), name='group_rule_delete'),
    url(r'^group/(?P<group_pk>[0-9]+)/rule/(?P<pk>[0-9]+)/update/$', GroupRuleUpdate.as_view(), name='group_rule_update'),

    url(r'^host/$', HostList.as_view(), name='host_list'),
    url(r'^host/add/$', HostCreate.as_view(), name='host_create'),
    url(r'^host/(?P<pk>[0-9]+)/$', HostDetail.as_view(), name='host_detail'),
    url(r'^host/(?P<pk>[0-9]+)/delete/$', HostDelete.as_view(), name='host_delete'),
    url(r'^host/(?P<pk>[0-9]+)/update/$', HostUpdate.as_view(), name='host_update'),

    url(r'^keyring/$', SSHKeyringList.as_view(), name='sshkeyring_list'),
    url(r'^keyring/add/$', SSHKeyringCreate.as_view(), name='sshkeyring_create'),
    url(r'^keyring/(?P<pk>[0-9]+)/$', SSHKeyringDetail.as_view(), name='sshkeyring_detail'),
    url(r'^keyring/(?P<pk>[0-9]+)/update/$', SSHKeyringUpdate.as_view(), name='sshkeyring_update'),
    url(r'^keyring/(?P<pk>[0-9]+)/delete/$', SSHKeyringDelete.as_view(), name='sshkeyring_delete'),

    url(r'^sshaccount/$', SSHAccountList.as_view(), name='sshaccount_list'),
    url(r'^sshaccount/(?P<pk>[0-9]+)/$', SSHAccountDetail.as_view(), name='sshaccount_detail'),
    url(r'^sshaccount/(?P<pk>[0-9]+)/delete/$', SSHAccountDelete.as_view(), name='sshaccount_delete'),
    url(r'^sshaccount/(?P<pk>[0-9]+)/update/$', SSHAccountUpdate.as_view(), name='sshaccount_update'),
    url(r'^sshaccount/(?P<pk>[0-9]+)/key/$', SSHAccountKeyUpdate.as_view(), name='sshaccount_key_update'),
    url(r'^sshaccount/add/environment$', SSHAccountCreateEnvironment.as_view(), name='sshaccount_create_environment'),
    url(r'^sshaccount/add/host$', SSHAccountCreateHost.as_view(), name='sshaccount_create_host'),
    url(r'^sshaccount/add/group$', SSHAccountCreateGroup.as_view(), name='sshaccount_create_group'),

    url(r'^sshaccountavailable/$', SSHAccountAvailableList.as_view(), name='sshaccountavailable_list'),
    url(r'^sshaccountavailable/(?P<pk>[0-9]+)/delete/$', SSHAccountAvailableDelete.as_view(), name='sshaccountavailable_delete'),

    url(r'^audit/key2access/$', AuditKey2Access.as_view(), name='audit_key2access'),

    url(r'^api/getkeys/$', api_get_keys),

    url(r'^$', HomeView.as_view(), name='home'),
)


import numpy as np
import math
import torch
from torch import autograd
from parameters import m, n, J, delta_t, delta_t_gen, H_design, B, C, B_mod, C_mod, delta_t_mod, J_mod, H_mod, H_design_inv, H_mod_inv, H_rotated, J_inacc

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

def f_gen(x):

    #A = torch.add(torch.einsum('nhw,wa->nh', B, x).T,C)
    A = torch.add(torch.reshape(torch.matmul(B, x),(m,m)).T,C)#.to(dev, non_blocking=True)

    # Taylor Expansion for F
    F = torch.eye(m)#.to(dev, non_blocking=True)
    for j in range(1,J+1):
        F_add = torch.matrix_power(A*delta_t_gen, j)/math.factorial(j)
        F = torch.add(F, F_add)

    return torch.matmul(F, x)

def f(x):

    #A = torch.add(torch.einsum('nhw,wa->nh', B, x).T,C)
    A = torch.add(torch.reshape(torch.matmul(B, x),(m,m)).T,C)#.to(dev, non_blocking=True)

    # Taylor Expansion for F
    F = torch.eye(m)#.to(dev, non_blocking=True)
    for j in range(1,J+1):
        F_add = torch.matrix_power(A*delta_t, j)/math.factorial(j)
        F = torch.add(F, F_add)

    return torch.matmul(F, x)

def f_approx(x):
    #A = torch.add(torch.einsum('nhw,wa->nh', B, x).T,C)
    A = torch.add(torch.reshape(torch.matmul(B, x),(m,m)).T,C)#.to(dev, non_blocking=True)

    # Taylor Expansion for F
    F = torch.eye(m)#.to(dev, non_blocking=True)
    for j in range(1,J_inacc+1):
        F_add = torch.matrix_power(A*delta_t, j)/math.factorial(j)
        F = torch.add(F, F_add)

    return torch.matmul(F, x)

def h(x):
    x = x#.to(dev, non_blocking=True)
    return torch.matmul(H_design,x)
    #return toSpherical(x)

def h_rotated(x):
    x = x#.to(dev, non_blocking=True)
    return torch.matmul(H_rotated,x)
    #return toSpherical(x)

def h_nonlinear(x):
    return toSpherical(x)

def f_interpolate(x, n=2):

    for _ in range(n):
        A = torch.add(torch.reshape(torch.matmul(B_mod, x),(m,m)).T,C_mod)#.to(dev, non_blocking=True)

        # Taylor Expansion for F
        F = torch.eye(m)#.to(dev, non_blocking=True)
        for j in range(1,J+1):
            F_add = torch.matrix_power(A*delta_t/n, j)/math.factorial(j)
            F = torch.add(F, F_add)
        x = torch.matmul(F, x)
        return x

def f_interpolate_approx(x, n=2):

    for _ in range(n):
        A = torch.add(torch.reshape(torch.matmul(B_mod, x),(m,m)).T,C_mod)#.to(dev, non_blocking=True)

        # Taylor Expansion for F
        F = torch.eye(m)#.to(dev, non_blocking=True)
        for j in range(1,J_inacc+1):
            F_add = torch.matrix_power(A*delta_t/n, j)/math.factorial(j)
            F = torch.add(F, F_add)
        x = torch.matmul(F, x)
    return x

def hInacc(x):
    return torch.matmul(H_mod,x)
    #return toSpherical(x)

"""
def getJacobian(x, a):

    if(x.size()[1] == 1):
        y = torch.reshape((x.T),[x.size()[0]])

    if(a == 'ObsAcc'):
        g = h
    elif(a == 'ModAcc'):
        g = f
    elif(a == 'ObsInacc'):
        g = h_rotated
    elif(a == 'ModInacc'):
        g = f_approx

    Jac = autograd.functional.jacobian(g, y)
    Jac = Jac.view(-1,m)
    return Jac

"""
def getJacobian(x, f):

    if(x.size()[1] == 1):
        y = torch.reshape((x.T),[x.size()[0]])

    Jac = autograd.functional.jacobian(f, y)
    Jac = Jac.view(-1,m)
    return Jac

def toSpherical(cart):

    rho = torch.norm(cart, p=2).view(1,1)
    phi = torch.atan2(cart[1, ...], cart[0, ...]).view(1, 1)
    phi = phi + (phi < 0).type_as(phi) * (2 * np.pi)

    theta = torch.acos(cart[2, ...] / rho).view(1, 1)

    spher = torch.cat([rho, theta, phi], dim=0)

    return spher

def toCartesian(sphe):

    rho = sphe[0,:]
    theta = sphe[1,:]
    phi = sphe[2,:]

    x = (rho * torch.sin(theta) * torch.cos(phi)).view(1,-1)
    y = (rho * torch.sin(theta) * torch.sin(phi)).view(1,-1)
    z = (rho * torch.cos(theta)).view(1,-1)

    cart = torch.cat([x,y,z],dim=0)

    return cart

def hInv(y):
    return torch.matmul(H_design_inv,y)
    #return toCartesian(y)

def hInaccInv(y):
    return torch.matmul(H_mod_inv,y)
    #return toCartesian(y)

'''
x = torch.tensor([[1],[1],[1]]).float()
H = getJacobian(x, 'ObsAcc')
print(H)
print(h(x))

F = getJacobian(x, 'ModAcc')
print(F)
print(f(x))
'''


import typing

from .command_composer import CommandComposer
from .command_reply_composer import CommandReplyComposer
from .event_composer import EventComposer
from .. import (
    Message,
    MessageType,
    MessageBusProtocol,
    Command,
    CommandType,
    CommandReply,
    CommandReplyType,
    Event,
    EventType,
)
from ....utils import UnitID

class MessageBuilder:
    """
    A helper class to easily build and emit new messages.

    This class stores a reference to the global message bus and offers methods to easily create new messages and send them through the bus.
    """

    def __init__(self, origin_id: UnitID, message_bus: MessageBusProtocol):
        """
        Args:
            origin_id: The component identifier of the origin of newly created messages.
            message_bus: The global message bus to use.
        """
        self._origin_id = origin_id
        self._message_bus = message_bus

        self._counters: typing.Dict[MessageType, int] = {
            CommandType: 0,
            CommandReplyType: 0,
            EventType: 0,
        }

    def build_command(
        self, cmd_type: type[CommandType], chain: Message | None = None, **kwargs
    ) -> CommandComposer:
        """
        Builds a new command.

        Args:
            cmd_type: The command type (i.e., a subclass of ``Command``).
            chain: A message that acts as the *predecessor* of the new message. Used to keep the same trace for multiple messages.
            **kwargs: Additional message parameters.

        Returns:
            A message composer to further build and finally emit the message.

        Raises:
            RuntimeError: ``cmd_type`` is not a subclass of ``Command``.
        """
        if not issubclass(cmd_type, Command):
            raise RuntimeError(f"Tried to build a command, but got a {cmd_type}")

        self._counters[CommandType] += 1

        return CommandComposer(
            self._origin_id, self._message_bus, cmd_type, chain, **kwargs
        )

    def build_command_reply(
        self,
        reply_type: type[CommandReplyType],
        command: CommandType,
        success: bool = True,
        message: str = "",
        **kwargs,
    ) -> CommandReplyComposer:
        """
        Builds a new command reply.

        Args:
            reply_type: The command reply type (i.e., a subclass of ``CommandReply``).
            command: The ``Command`` this reply is based on.
            success: Whether the command *succeeded* or not (how this is interpreted depends on the command).
            message: Additional message to include in the reply.
            **kwargs: Additional message parameters.

        Returns:
            A message composer to further build and finally emit the message.

        Raises:
            RuntimeError: ``reply_type`` is not a subclass of ``CommandReply``.
        """
        if not issubclass(reply_type, CommandReply):
            raise RuntimeError(
                f"Tried to build a command reply, but got a {reply_type}"
            )

        self._counters[CommandReplyType] += 1

        return CommandReplyComposer(
            self._origin_id,
            self._message_bus,
            reply_type,
            command,
            success=success,
            message=message,
            unique=command.unique,
            **kwargs,
        )

    def build_event(
        self, event_type: type[Event], chain: Message | None = None, **kwargs
    ) -> EventComposer:
        """
        Builds a new event.

        Args:
            event_type: The event type (i.e., a subclass of ``Event``).
            chain: A message that acts as the *predecessor* of the new message. Used to keep the same trace for multiple messages.
            **kwargs: Additional message parameters.

        Returns:
            A message composer to further build and finally emit the message.

        Raises:
            RuntimeError: ``event_type`` is not a subclass of ``Event``.
        """
        if not issubclass(event_type, Event):
            raise RuntimeError(f"Tried to build an event, but got a {event_type}")

        self._counters[EventType] += 1

        return EventComposer(
            self._origin_id, self._message_bus, event_type, chain, **kwargs
        )

    def get_message_count(self, msg_type: MessageType) -> int:
        """
        Gets how many messages of specific message types have been created.

        The message builder keeps track of how many messages of the three major types ``Command``, ``CommandReply`` and
        ``Event`` have been created.

        Args:
            msg_type: The message type to get the count of. Must be either ``Command``, ``CommandReply`` or ``Event``.

        Returns:
            The number of messages already built of the specified type.
        """
        return self._counters[msg_type] if msg_type in self._counters else 0


import torch
import time
from AttentionModel import BigramLanguageModel

#hyperparams
batch_size = 4      #64 #how many independent sequences will we process in parallel
block_size = 8       #256 #what is the maximum context length for predictions
max_iters = 5000
learning_rate = 1e-3      #3e-4
eval_interval = 500
eval_iters = 200
n_embed = 32       #384
num_heads = 4 #6
n_layer = 3 #how many decoder blocks in the model
dropout = 0.0
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - block_size, (batch_size, ))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    x, y = x.to(device), y.to(device)
    return x, y

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

if __name__ == '__main__':
    with open('input.txt', 'r', encoding='utf-8') as f:
        text = f.read()
        print("length of dataset (chars)", len(text))
        print(type(text))
        all_chars = sorted(list(set(text)))

    ctoi = {ch:i for i, ch in enumerate(all_chars)}
    itoc = {i:ch for i, ch in enumerate(all_chars)}

    vocab_size = len(all_chars)
    encode = lambda s: [ctoi[c] for c in s]
    decode = lambda i: ''.join([itoc[c] for c in i])

    print(encode("I be busy in the street I be riding on his meat"))
    print(decode(encode("I be busy in the street I be riding on his meat")))
    data = torch.tensor(encode(text), dtype=torch.long)
    print(data.shape, data.dtype)

    n = int(0.9*len(data))
    train_data = data[:n]
    val_data = data[n:]

    torch.manual_seed(1337)

    print(device)

    xb, yb = get_batch('train')

    model = BigramLanguageModel(vocab_size, n_embed, block_size, num_heads, n_layer, dropout)
    m = model.to(device)
    optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)

    startTime = time.time()
    for step in range(max_iters):
        xb, yb = get_batch('train')

        logits, loss = m(xb, yb)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

        if step % eval_interval == 0:
            losses = estimate_loss()
            print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")

    print('training time:', time.time() - startTime)
    context = torch.zeros((1, 1), dtype=torch.long, device=device)
    print(decode(m.generate(context, max_new_tokens=100)[0].tolist()))