import numpy as np
import torch as tr
import torch.nn as nn

device = tr.device("cuda") if tr.cuda.is_available() else tr.device("cpu")

def printDataStats(data):
    nUsers = len(data["users"])
    print("User features: %s" % data["users"][0].shape)
    print("Session features: %d" % (data["sessions"][0][0].shape))
    print("Hits features: %d" % (data["hits"][0][0][0].shape))
    print("Targets features: %d" % (data["targets"][0][0].shape))

    for i in range(nUsers):
        nSessions = len(data["sessions"][i])
        print("User %d. Num sessions: %d" % (i, nSessions))
        for j in range(nSessions):
            nHits = len(data["hits"][j])
            print("  - session %d. Num hits: %d" % (j, nHits))

def getData(nUsers, NU, NS, NH, NT):
    dataUsers = tr.randn(nUsers, NU).to(device)
    dataSessions = []
    dataHits = []
    dataTargets = []
    # For each user, generate some number of sessions between 1 and 10
    for i in range(nUsers):
        sessCount = np.random.randint(1, 10)
        dataSessions.append(tr.randn(sessCount, NS).to(device))
        dataTargets.append(tr.randn(sessCount, NT).to(device))
        # For each session, generate some number of hits between 1 and 10
        sessionHits = []
        for j in range(sessCount):
            hitsCount = np.random.randint(1, 10)
            sessionHits.append(tr.randn(hitsCount, NH).to(device))
        dataHits.append(sessionHits)

    data = {"users" : dataUsers, "sessions" : dataSessions, "hits" : dataHits, "targets" : dataTargets}
    return data

class Model(nn.Module):
    def __init__(self, NU, NS, NH, NT):
        super(Model, self).__init__()
        self.rnnHits = nn.RNN(input_size=NH, hidden_size=10)
        self.rnnSessions = nn.RNN(input_size=NS, hidden_size=10)
        self.linearOut = nn.Linear(10 + 10 + NU, NT)

    def forward(self, data):
        # Hidden state for each session of each user:
        #  hiddenSessions :: [numSessions x 1 (batch) x 10] for each user (total nUsers)
        hiddenSessions = [self.rnnSessions(data["sessions"][i].unsqueeze(dim=1))[0][:, 0] \
            for i in range(len(data["sessions"]))]

        # Hidden state for each hits of each session of each user:
        hiddenHits = []
        for i in range(len(data["sessions"])):
            userNumSessions = len(data["sessions"][i])
            # Hidden state for each session of ith user user:
            #  userHits :: [numHits x 1 (batch) x 10] for this user (total userNumSessions)
            userHits =  [self.rnnHits(data["hits"][i][j].unsqueeze(dim=1))[0] \
                for j in range(len(data["hits"][i]))]
            # Get only the last hidden state (basically the cummulation of all hidden states), so we transform the
            #  hits into session features.
            lastUserHit = tr.cat([x[0] for x in userHits]).to(device)
            hiddenHits.append(lastUserHit)

        # Now, concatenate everything, resulting a list of nUsers x [nSessionsEachUser, NU + 10 + 10]
        allFeatures = []
        for i in range(len(data["users"])):
            nSessions = len(data["sessions"][i])
            user = data["users"][i].unsqueeze(dim=0).repeat(nSessions, 1)
            session = hiddenSessions[i]
            hit = hiddenHits[i]
            concatenated = tr.cat([user, session, hit], dim=1)
            allFeatures.append(concatenated)

        # Finally, project this output to the targets space
        results = [self.linearOut(x) for x in allFeatures]
        return results

def main():
    nUsers = 10
    # number of features per users, sessions and hits and target valiues (5 shopping stages)
    NU, NS, NH, NT = 5, 10, 15, 5
    data = getData(nUsers, NU, NS, NH, NT)

    # Print for each user the generated data
    printDataStats(data)

    # Pass the data through the users>sessions>hits model
    model = Model(NU, NS, NH, NT).to(device)
    res = model.forward(data)
    print([x.shape for x in res])
    print([x.shape for x in data["targets"]])

    # Compute the loss (note, in reality we use one hot encoding, not [0-4] values for targets and cross entropy)
    loss = [((x - y)**2).mean() for x, y in zip(res, data["targets"])]
    loss = sum(loss)
    print("Loss: %2.5f" % loss)

    # Optimize the loss, use train-val split etc.
    # ...

if __name__ == "__main__":
    main()