#!/usr/bin/env python

from __future__ import print_function

import os
import sys
import csv
import numpy as np
import math
import random
from collections import defaultdict
import torch
from torch.autograd import Variable
from torch.nn.parameter import Parameter
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

sys.path.append( '%s/../prog_common' % os.path.dirname(os.path.realpath(__file__)) )
from prog_util import DECISION_DIM
from cmd_args import cmd_args
from pytorch_initializer import weights_init

sys.path.append( '%s/../cfg_parser' % os.path.dirname(os.path.realpath(__file__)) )
import cfg_parser as parser

class CNNEncoder(nn.Module):
    def __init__(self, max_len, latent_dim):
        super(CNNEncoder, self).__init__()
        self.latent_dim = latent_dim
        self.max_len = max_len

        self.conv1 = nn.Conv1d(DECISION_DIM, cmd_args.c1, cmd_args.c1)
        self.conv2 = nn.Conv1d(cmd_args.c1, cmd_args.c2, cmd_args.c2)
        self.conv3 = nn.Conv1d(cmd_args.c2, cmd_args.c3, cmd_args.c3)

        self.last_conv_size = max_len - cmd_args.c1 + 1 - cmd_args.c2 + 1 - cmd_args.c3 + 1
        self.w1 = nn.Linear(self.last_conv_size * cmd_args.c3, cmd_args.dense)
        self.mean_w = nn.Linear(cmd_args.dense, latent_dim)
        self.log_var_w = nn.Linear(cmd_args.dense, latent_dim)
        weights_init(self)

    def forward(self, x_cpu):
        if cmd_args.mode == 'cpu':
            batch_input = Variable(torch.from_numpy(x_cpu))
        else:
            batch_input = Variable(torch.from_numpy(x_cpu).cuda())

        h1 = self.conv1(batch_input)
        h1 = F.relu(h1)        
        h2 = self.conv2(h1)
        h2 = F.relu(h2)
        h3 = self.conv3(h2)
        h3 = F.relu(h3)

        # h3 = torch.transpose(h3, 1, 2).contiguous()
        flatten = h3.view(x_cpu.shape[0], -1)
        h = self.w1(flatten)
        h = F.relu(h)

        z_mean = self.mean_w(h)
        z_log_var = self.log_var_w(h)
        
        return (z_mean, z_log_var)

if __name__ == '__main__':
    pass

[((33, 21, 33, 70), 'torch.nn.Conv1d', 'nn.Conv1d', ({(33, 31, 33, 43): 'DECISION_DIM', (33, 45, 33, 56): 'cmd_args.c1', (33, 58, 33, 69): 'cmd_args.c1'}, {}), '(DECISION_DIM, cmd_args.c1, cmd_args.c1)', True, 'import torch.nn as nn\n'), ((34, 21, 34, 69), 'torch.nn.Conv1d', 'nn.Conv1d', ({(34, 31, 34, 42): 'cmd_args.c1', (34, 44, 34, 55): 'cmd_args.c2', (34, 57, 34, 68): 'cmd_args.c2'}, {}), '(cmd_args.c1, cmd_args.c2, cmd_args.c2)', True, 'import torch.nn as nn\n'), ((35, 21, 35, 69), 'torch.nn.Conv1d', 'nn.Conv1d', ({(35, 31, 35, 42): 'cmd_args.c2', (35, 44, 35, 55): 'cmd_args.c3', (35, 57, 35, 68): 'cmd_args.c3'}, {}), '(cmd_args.c2, cmd_args.c3, cmd_args.c3)', True, 'import torch.nn as nn\n'), ((38, 18, 38, 78), 'torch.nn.Linear', 'nn.Linear', ({(38, 28, 38, 61): 'self.last_conv_size * cmd_args.c3', (38, 63, 38, 77): 'cmd_args.dense'}, {}), '(self.last_conv_size * cmd_args.c3, cmd_args.dense)', True, 'import torch.nn as nn\n'), ((39, 22, 39, 59), 'torch.nn.Linear', 'nn.Linear', ({(39, 32, 39, 46): 'cmd_args.dense', (39, 48, 39, 58): 'latent_dim'}, {}), '(cmd_args.dense, latent_dim)', True, 'import torch.nn as nn\n'), ((40, 25, 40, 62), 'torch.nn.Linear', 'nn.Linear', ({(40, 35, 40, 49): 'cmd_args.dense', (40, 51, 40, 61): 'latent_dim'}, {}), '(cmd_args.dense, latent_dim)', True, 'import torch.nn as nn\n'), ((41, 8, 41, 26), 'pytorch_initializer.weights_init', 'weights_init', ({(41, 21, 41, 25): 'self'}, {}), '(self)', False, 'from pytorch_initializer import weights_init\n'), ((50, 13, 50, 23), 'torch.nn.functional.relu', 'F.relu', ({(50, 20, 50, 22): 'h1'}, {}), '(h1)', True, 'import torch.nn.functional as F\n'), ((52, 13, 52, 23), 'torch.nn.functional.relu', 'F.relu', ({(52, 20, 52, 22): 'h2'}, {}), '(h2)', True, 'import torch.nn.functional as F\n'), ((54, 13, 54, 23), 'torch.nn.functional.relu', 'F.relu', ({(54, 20, 54, 22): 'h3'}, {}), '(h3)', True, 'import torch.nn.functional as F\n'), ((59, 12, 59, 21), 'torch.nn.functional.relu', 'F.relu', ({(59, 19, 59, 20): 'h'}, {}), '(h)', True, 'import torch.nn.functional as F\n'), ((19, 55, 19, 81), 'os.path.realpath', 'os.path.realpath', ({(19, 72, 19, 80): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((24, 54, 24, 80), 'os.path.realpath', 'os.path.realpath', ({(24, 71, 24, 79): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((45, 35, 45, 58), 'torch.from_numpy', 'torch.from_numpy', ({(45, 52, 45, 57): 'x_cpu'}, {}), '(x_cpu)', False, 'import torch\n'), ((47, 35, 47, 58), 'torch.from_numpy', 'torch.from_numpy', ({(47, 52, 47, 57): 'x_cpu'}, {}), '(x_cpu)', False, 'import torch\n')]