# Parse local options first, and rewrite the sys.argv[].
# We need to do that before import "common", as otherwise we get an error for
# unrecognized arguments.
import argparse
import os
import sys

parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--replicated', action='store_true')
parser.add_argument('--long_test', action='store_true')
parser.add_argument('--max_diff_count', type=int, default=25)
FLAGS, leftovers = parser.parse_known_args()
sys.argv = [sys.argv[0]] + leftovers
# Setup import folders.
_XLA_FOLDER = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
sys.path.append(os.path.join(os.path.dirname(_XLA_FOLDER), 'test'))
sys.path.insert(0, _XLA_FOLDER)

# Normal imports section starts here.
import collections
from common_utils import TestCase, run_tests, iter_indices
import itertools
import numpy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch_xla
import torch_xla_py.utils as xu
import torch_xla_py.xla_model as xm
import torchvision
import unittest

DeviceSupport = collections.namedtuple('DeviceSupport', ['num_devices'])

def _gen_tensor(*args, **kwargs):
  return torch.randn(*args, **kwargs)

def _gen_int_tensor(*args, **kwargs):
  return torch.randint(*args, **kwargs)

class Holder(object):
  pass

def _get_device_support(devname):
  assert devname in ['TPU', 'CPU']
  # If the Cloud TPU config file is present, we support TPUs.
  if (os.path.isfile(os.path.join(os.environ['HOME'], '.pytorch_tpu.conf')) or
      os.environ.get('XRT_TPU_CONFIG', None)):
    if devname == 'TPU':
      return DeviceSupport(
          num_devices=int(os.environ.get('TPU_NUM_DEVICES', 8)))
    return DeviceSupport(num_devices=int(os.environ.get('CPU_NUM_DEVICES', 1)))

  xrt_devmap = os.environ.get('XRT_DEVICE_MAP', None)
  if xrt_devmap is None:
    return None
  num_devices = 0
  for dev_spec in xrt_devmap.split('|'):
    dev_parts = dev_spec.split(';')
    if dev_parts[0].startswith(devname):
      num_devices += 1
  if num_devices > 0:
    return DeviceSupport(num_devices=num_devices)

def _support_replicated(devname, num_devices):
  devsup = _get_device_support(devname)
  if not devsup:
    return False
  return devsup.num_devices >= num_devices

def _random_inputs(shapes, num_replicas=1):
  random_tensors = []
  for _ in range(0, num_replicas):
    replica_inputs = []
    for shape in shapes:
      replica_inputs.append(_gen_tensor(*shape))
    random_tensors.append(tuple(replica_inputs))
  return tuple(random_tensors)

def _random_like(tensor_list):
  random_tensors = []
  for o in tensor_list:
    if o.dtype == torch.float32 or o.dtype == torch.float64:
      random_tensors += [_gen_tensor(*o.shape, dtype=o.dtype)]
    elif o.dtype == torch.int64:
      # TODO remove this, we shouldn't be needing to pass random_tensor for long types
      random_tensors += [torch.empty_like(o)]
    else:
      raise RuntimeError('Unsupported type: ', o.dtype)
  return random_tensors

def _zeros_like(tensor_list):
  zeros_tensors = []
  for o in tensor_list:
    if o.dtype == torch.float32 or o.dtype == torch.float64:
      zeros_tensors += [torch.zeros(*o.shape, dtype=o.dtype)]
    elif o.dtype == torch.int64:
      # TODO remove this, we shouldn't be needing to pass zeros_tensor for long types
      zeros_tensors += [torch.zeros_like(o)]
    else:
      raise RuntimeError('Unsupported type: ', o.dtype)
  return zeros_tensors

def _dump_differences(target, result, rtol=1e-5, atol=1e-3, max_diff_count=0):
  env = Holder()
  env.max_diff = 0.0
  env.max_rel = None
  env.max_index = None
  env.diff_count = 0

  def check_values(a, b, index):
    r = max(abs(a), abs(b)) * rtol
    diff = abs(a - b)
    if diff > max(r, atol):
      print('a={}\tb={}\tdiff={}\tindex={}'.format(a, b, diff, index))
      env.diff_count += 1
      if diff > env.max_diff:
        env.max_diff = diff
        env.max_rel = diff / max(abs(a), abs(b))
        env.max_index = index

  if isinstance(target, torch.Tensor):
    assert isinstance(result, torch.Tensor)
    assert target.size() == result.size()
    if target.dim() > 0:
      for i in iter_indices(target):
        check_values(target[i], result[i], i)
        if max_diff_count > 0 and env.diff_count >= max_diff_count:
          break
    else:
      check_values(target.item(), result.item(), 0)
  elif isinstance(target, (list, tuple)):
    assert isinstance(result, (list, tuple))
    assert len(target) == len(result)
    for i, v in enumerate(target):
      check_values(v, result[i], [i])
      if max_diff_count > 0 and env.diff_count >= max_diff_count:
        break
  elif isinstance(target, float):
    assert isinstance(result, float)
    check_values(target, result, [])
  if env.max_index is not None:
    print('\nmax_diff={}\tmax_rel={}\tindex={}'.format(
        env.max_diff, env.max_rel, env.max_index))

def _xla_run(model, input, device='TPU'):
  if isinstance(input, (tuple, list)):
    devices = ['{}:{}'.format(device, n) for n in range(0, len(input))]
    xla_model = xm.XlaModel(
        model,
        input[0],
        num_cores=len(input),
        devices=devices,
        full_conv_precision=True)
    output_xla = xla_model(*input)
    return xm.convert_to_tensors(output_xla)
  else:
    xla_model = xm.XlaModel(model, [input], full_conv_precision=True)
    output_xla = xla_model(input)
    return output_xla[0]

class XlaTestCase(TestCase):

  def assertEqualRel(self, out, expected, rel_err=1e-2, abs_err=1e-5):
    try:
      diff_tensor = (out - expected).abs()
      max_rel_err = torch.max(out.abs(), expected.abs()) * rel_err
      # Allow higher relative differences as long as we're still below the
      # absolute error.
      max_abs_err = torch.max(max_rel_err, torch.ones_like(out) * abs_err)
      super(XlaTestCase, self).assertEqual(diff_tensor.size(),
                                           max_abs_err.size())
      if torch.le(diff_tensor, max_abs_err).min().item() == 0:
        self.fail('Relative error higher than the maximum tolerance')
    except:
      _dump_differences(
          out,
          expected,
          rtol=rel_err,
          atol=abs_err,
          max_diff_count=FLAGS.max_diff_count)
      raise

  def assertEqualDbg(self, out, expected):
    try:
      super(XlaTestCase, self).assertEqual(out, expected)
    except:
      _dump_differences(
          out,
          expected,
          rtol=1e-8,
          atol=1e-8,
          max_diff_count=FLAGS.max_diff_count)
      raise

  def compareReplicated(self, model, inputs, xla_outputs):
    self.assertEqual(len(inputs), len(xla_outputs))
    for i, input in enumerate(inputs):
      expected = xu.as_list(model(*input))
      xla_output = xu.as_list(xla_outputs[i])
      self.assertEqual(len(expected), len(xla_output))
      for j, expected_tensor in enumerate(expected):
        self.assertEqualDbg(xla_output[j], expected_tensor)

  def compareModel(self, model, input, rel_err=0.05, abs_err=1e-4):
    xla_model = xm.XlaModel(model, [input], full_conv_precision=True)
    output_xla = xla_model(input)
    output = model(input)
    self.assertEqualRel(
        output,
        xm.convert_to_tensors(output_xla)[0],
        rel_err=rel_err,
        abs_err=abs_err)
    grad_output = _gen_tensor(*output.shape)  # random gradients
    grad_output.grad = grad_output.data
    output.backward(grad_output)
    xla_model.backward([grad_output])
    xla_updated_params = [p.grad.to_tensor() for p in xla_model.parameters()[0]]
    updated_params = [p.grad for p in model.parameters()]
    self.assertEqual(len(xla_updated_params), len(updated_params))
    for i in range(0, len(updated_params)):
      self.assertEqualRel(
          xla_updated_params[i],
          updated_params[i],
          rel_err=rel_err,
          abs_err=abs_err)

class TestMulAdd(XlaTestCase):

  def test(self):

    class XlaMulAdd(nn.Module):

      def forward(self, x, y):
        return x * y + y

    x = torch.rand(3, 5)
    y = torch.rand(3, 5)
    model = XlaMulAdd()
    traced_model = torch.jit.trace(model, (x, y))
    xla_model = torch_xla._XLAC.XlaModule(traced_model)
    inputs_xla = [torch_xla._XLAC.XLATensor(x), torch_xla._XLAC.XLATensor(y)]
    output_xla = xla_model((tuple(inputs_xla)))
    expected = model(x, y)
    self.assertEqualDbg(output_xla[0][0].to_tensor().data, expected.data)

class TestRelu(XlaTestCase):

  def test(self):

    class XlaRelu(nn.Module):

      def forward(self, x):
        return F.relu(x)

    x = _gen_tensor(2, 1, 4, 6)
    model = XlaRelu()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestThreshold(XlaTestCase):

  def test(self):

    class XlaThreshold(nn.Module):

      def __init__(self):
        super(XlaThreshold, self).__init__()
        self.threshold = nn.Threshold(0.4, 20)

      def forward(self, x):
        return self.threshold(x)

    x = torch.rand(4, 2)
    model = XlaThreshold()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestTranspose(XlaTestCase):

  def test(self):

    class XlaTranspose(nn.Module):

      def forward(self, x):
        return torch.t(x)

    x = torch.rand(2, 3)
    model = XlaTranspose()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestView(XlaTestCase):

  def test(self):

    class XlaView(nn.Module):

      def forward(self, x):
        return x.view(-1, 16)

    x = torch.rand(4, 8)
    model = XlaView()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestStack(XlaTestCase):

  def test(self):

    class XlaStack(nn.Module):

      def __init__(self, dim):
        super(XlaStack, self).__init__()
        self.dim = dim

      def forward(self, x, y):
        return torch.stack((x, y), self.dim)

    x = torch.rand(2, 5)
    y = torch.rand(2, 5)
    for dim in [0, 1]:
      model = XlaStack(dim)
      traced_model = torch.jit.trace(model, (x, y))
      xla_model = torch_xla._XLAC.XlaModule(traced_model, differentiate=False)
      inputs_xla = [torch_xla._XLAC.XLATensor(x), torch_xla._XLAC.XLATensor(y)]
      output_xla = xla_model((tuple(inputs_xla)))
      expected = model(x, y)
      self.assertEqualDbg(output_xla[0][0].to_tensor().data, expected.data)

class TestExpand(XlaTestCase):

  def test(self):

    class XlaExpand(nn.Module):

      def forward(self, x):
        return x.expand(2, 5)

    x = torch.rand(5)
    model = XlaExpand()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestLinear(XlaTestCase):

  def test(self):

    class XlaLinear(nn.Module):

      def __init__(self):
        super(XlaLinear, self).__init__()
        self.linear = nn.Linear(2, 5)

      def forward(self, x):
        return self.linear(x)

    x = torch.rand(4, 2)
    model = XlaLinear()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestNonContiguousTensor(XlaTestCase):

  def test(self):

    class XlaPlusSelf(nn.Module):

      def forward(self, x):
        return x + x

    x = torch.rand(3, 7)
    model = XlaPlusSelf()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)
    out_t = _xla_run(model, x.t())
    expected_t = model(x.t())
    self.assertEqualDbg(out_t.data, expected_t.data)
    self.assertEqualDbg(out_t.data, out.t().data)

class TestConstantTensor(XlaTestCase):

  def test(self):
    x = torch.rand(2, 2)
    y = torch.rand(2, 2)

    class XplusY(nn.Module):

      def forward(self, a):
        return a + y

    class XmulY(nn.Module):

      def forward(self, a):
        return a * y

    model = XplusY()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

    model = XmulY()
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class TestConv(XlaTestCase):

  def test(self):

    class XlaConv(nn.Module):

      def __init__(self, stride, padding, bias):
        super(XlaConv, self).__init__()
        self.conv = nn.Conv2d(
            10, 100, 5, stride=stride, padding=padding, bias=bias)

      def forward(self, x):
        return self.conv(x)

    for stride in range(1, 4):
      for padding in range(0, 3):
        for bias in [True, False]:
          x = _gen_tensor(32, 10, 28, 28)
          model = XlaConv(stride, padding, bias)
          out = _xla_run(model, x)
          expected = model(x)
          self.assertEqualRel(out.data, expected.data)

class TestMaxPool(XlaTestCase):

  def test(self):

    class XlaMaxPool(nn.Module):

      def __init__(self, stride, padding):
        super(XlaMaxPool, self).__init__()
        self.stride = stride
        self.padding = padding

      def forward(self, x):
        return F.max_pool2d(x, 3, stride=self.stride, padding=self.padding)

    x = torch.rand(1, 64, 112, 112)
    for stride in [None, 2]:
      for padding in [0, 1]:
        model = XlaMaxPool(stride, padding)
        out = _xla_run(model, x)
        expected = model(x)
        self.assertEqualDbg(out.data, expected.data)

class TestAvgPool(XlaTestCase):

  def test(self):

    class XlaAvgPool(nn.Module):

      def __init__(self, stride, padding, count_include_pad):
        super(XlaAvgPool, self).__init__()
        self.stride = stride
        self.padding = padding
        self.count_include_pad = count_include_pad

      def forward(self, x):
        return F.avg_pool2d(x, 2, self.stride, self.padding, False,
                            self.count_include_pad)

    x = torch.rand(1, 1, 3, 3)
    for stride in [1, 2, None]:
      for padding in [0, 1]:
        for count_include_pad in [False, True]:
          model = XlaAvgPool(stride, padding, count_include_pad)
          out = _xla_run(model, x)
          expected = model(x)
          self.assertEqualDbg(out.data, expected.data)

class TestLogSoftmax(XlaTestCase):

  def test(self):

    class XlaLogSoftmax(nn.Module):

      def __init__(self, dim):
        super(XlaLogSoftmax, self).__init__()
        self.dim = dim

      def forward(self, x):
        return F.log_softmax(x, self.dim)

    x = torch.rand(5, 3, 4, 2)
    for dim in range(0, x.dim()):
      model = XlaLogSoftmax(dim)
      out = _xla_run(model, x)
      expected = model(x)
      self.assertEqualRel(out.data, expected.data, rel_err=1e-4, abs_err=1)

class TestBatchNorm(XlaTestCase):

  def test(self):

    class XlaBatchNorm(nn.Module):

      def __init__(self, training):
        super(XlaBatchNorm, self).__init__()
        if training:
          self.bn = nn.BatchNorm2d(3)
        else:
          self.bn = nn.BatchNorm2d(3, track_running_stats=False)

      def forward(self, x):
        return self.bn(x)

    x = torch.rand(14, 3, 5, 7)
    model = XlaBatchNorm(True)
    out = _xla_run(model, x)
    expected = model(x)
    self.assertEqualDbg(out.data, expected.data)

class XlaMNIST(nn.Module):

  def __init__(self):
    super(XlaMNIST, self).__init__()
    self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
    self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
    self.fc1 = nn.Linear(320, 50)
    self.fc2 = nn.Linear(50, 10)

  def forward(self, x):
    x = F.relu(F.max_pool2d(self.conv1(x), 2))
    x = F.relu(F.max_pool2d(self.conv2(x), 2))
    x = x.view(-1, 320)
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return F.log_softmax(x, dim=1)

class TestMNIST(XlaTestCase):

  def test(self):
    batch_size = 32
    x = _gen_tensor(batch_size, 1, 28, 28)
    model = XlaMNIST()
    self.compareModel(model, x)

class AxPlusB(nn.Module):

  def __init__(self, dims=None):
    super(AxPlusB, self).__init__()
    self.ones = torch.ones(*dims)
    self.a = nn.Parameter(_gen_tensor(1, 1))
    self.b = nn.Parameter(_gen_tensor(1, 1))

  def forward(self, x):
    return x.mm(self.a) + self.ones.mm(self.b)

class SquareLoss(nn.Module):

  def __init__(self):
    super(SquareLoss, self).__init__()

  def forward(self, x, y):
    x.requires_grad = True
    y.requires_grad = True
    diff = x - y
    loss = diff.t().mm(diff)[0][0]
    return loss / x.size()[0]

class TestAxPlusB(XlaTestCase):

  def test(self):
    A = 3.11
    B = 4.09
    model = AxPlusB(dims=(1, 1))
    xla_model = xm.XlaModel(model, [_gen_tensor(1, 1)])
    optimizer = optim.SGD(xla_model.parameters_list(), lr=0.1, momentum=0.5)
    square_loss = SquareLoss()
    loss = None
    for _ in range(0, 100):
      optimizer.zero_grad()
      x = _gen_tensor(1, 1)
      target = x * A + B
      y = xla_model(x)
      loss = square_loss(y[0], target)
      loss.backward()
      xla_model.backward(y)
      optimizer.step()
    self.assertEqualRel(loss.sum(), torch.tensor(0.0))

class TestAxPlusBGen(XlaTestCase):

  def test(self):
    A = 3.11
    B = 4.09
    batch_size = 128
    gen = xu.FnDataGenerator(
        lambda x: x * A + B, batch_size, _gen_tensor, count=100)
    model = AxPlusB(dims=(batch_size, 1))
    xla_model = xm.XlaModel(model, [_gen_tensor(batch_size, 1)])
    optimizer = optim.SGD(xla_model.parameters_list(), lr=0.1, momentum=0.5)
    square_loss = SquareLoss()
    loss = None
    for x, target in gen:
      optimizer.zero_grad()
      y = xla_model(x)
      loss = square_loss(y[0], target)
      loss.backward()
      xla_model.backward(y)
      optimizer.step()
    self.assertEqualRel(loss.sum(), torch.tensor(0.0))

class TestAxPlusBGenXla(XlaTestCase):

  def test(self):
    batch_size = 128
    scaler = torch.Tensor([[1.0 / batch_size]])

    def loss_fn(x, y):
      diff = x - y
      sloss = diff.t().mm(diff)
      return sloss.mm(scaler)

    A = 3.11
    B = 4.09
    gen = xu.FnDataGenerator(
        lambda x: x * A + B, batch_size, _gen_tensor, count=100)
    model = AxPlusB(dims=(batch_size, 1))
    xla_model = xm.XlaModel(
        model, [_gen_tensor(batch_size, 1)],
        target=_gen_tensor(batch_size, 1),
        loss_fn=loss_fn,
        num_cores=1,
        devices=[':0'])
    optimizer = optim.SGD(xla_model.parameters_list(), lr=0.1, momentum=0.5)
    xla_model.train(gen, optimizer, batch_size, log_fn=None)

    def eval_fn(output, target):
      mloss = (output - target) * (output - target)
      error = torch.ones_like(mloss) * 1e-5
      count = torch.le(mloss, error).sum()
      return mloss.mean().item(), count.item()

    gen = xu.FnDataGenerator(lambda x: x * A + B, batch_size, _gen_tensor)
    accuracy = xla_model.test(gen, eval_fn, batch_size, log_fn=None)
    self.assertEqual(accuracy, 100.0)

class TestCompareAxPlusB(XlaTestCase):

  def test(self):
    batch_size = 128
    model = AxPlusB(dims=(batch_size, 1))
    x = _gen_tensor(batch_size, 1)
    self.compareModel(model, x)

class TestSum(XlaTestCase):

  def test(self):

    class XlaSum(nn.Module):

      def __init__(self, dim):
        super(XlaSum, self).__init__()
        self.dim = dim

      def forward(self, x):
        return x.sum(dim=self.dim)

    x = _gen_tensor(2, 3, 4, 6)
    for dim in range(0, x.dim()):
      model = XlaSum(dim)
      out = _xla_run(model, x)
      expected = model(x)
      self.assertEqualDbg(out.data, expected.data)

class XlaNllLoss(nn.Module):

  def __init__(self):
    super(XlaNllLoss, self).__init__()
    self.nll_loss = nn.NLLLoss()

  def forward(self, x, labels):
    return self.nll_loss(x, labels)

class TestNllLoss(TestCase):

  def test(self):
    input = _gen_tensor(3, 5, requires_grad=True)
    target = torch.empty(3, dtype=torch.long).random_(5)
    model = XlaNllLoss()
    traced_model = torch.jit.trace(model, (input, target))
    xla_model = torch_xla._XLAC.XlaModule(traced_model)
    xla_inputs = [
        torch_xla._XLAC.XLATensor(input),
        torch_xla._XLAC.XLATensor(target)
    ]
    output_xla = xla_model((tuple(xla_inputs)))
    expected = model(input, target)
    self.assertEqual(output_xla[0][0].to_tensor().data, expected.data)

class TestLongGraphChain(XlaTestCase):

  def test(self):
    orig_x = torch.Tensor([[1, 2], [3, 4]])
    orig_y = torch.Tensor([[0.1, 0.2], [0.3, 0.4]])
    x = orig_x
    y = orig_y
    xla_x = torch_xla._XLAC.XLATensor(orig_x)
    xla_y = torch_xla._XLAC.XLATensor(orig_y)
    for i in range(0, 10000):
      x = x + 2 * y
      xla_x = xla_x.add(2, xla_y)
    self.assertEqualRel(x, xla_x.to_tensor(), rel_err=1e-3, abs_err=5)

class TestGradients(XlaTestCase):

  def checkGrad(self,
                model,
                inputs,
                grad_outputs='random',
                xla=True,
                rel_err=1e-2,
                abs_err=1e-5):
    # Trace and symbolically differentiate
    traced_model = torch.jit.trace(model, *inputs)
    fwd = traced_model._get_method('forward')
    xm.forward_passes(fwd.graph)

    inputs_params = inputs + list(model.parameters())
    inputs_params_buffers = inputs + list(fwd.params())

    gradient = torch._C._jit_differentiate(fwd.graph)
    xm.forward_passes(gradient.f)
    xm.backward_passes(gradient.df)

    ##############################################################
    # Run forward and backwarg graphs via jit interpreter
    exec_f = torch._C.GraphExecutor(gradient.f, False)
    exec_df = torch._C.GraphExecutor(gradient.df, False)

    # forward function
    raw_outputs = exec_f(*inputs_params_buffers)
    raw_outputs = xu.as_list(raw_outputs)
    intermediate_outputs = [
        raw_output for raw_output in raw_outputs[gradient.f_real_outputs:]
        if isinstance(raw_output, torch.Tensor)
    ]
    outputs = raw_outputs[:gradient.f_real_outputs]

    if grad_outputs == 'random':
      grad_outputs = _random_like(outputs) + _zeros_like(intermediate_outputs)

    raw_grad_outputs = []
    raw_grad_outputs += grad_outputs
    raw_grad_outputs += [
        inputs_params_buffers[i] for i in gradient.df_input_captured_inputs
    ]
    raw_grad_outputs += [
        raw_outputs[i] for i in gradient.df_input_captured_outputs
    ]

    grad_inputs = exec_df(*raw_grad_outputs)
    grad_inputs = xu.as_list(grad_inputs)

    ##############################################################
    # backward with XLA
    if xla:
      xla_model = torch_xla._XLAC.XlaModule(traced_model)
      inputs_xla = [torch_xla._XLAC.XLATensor(input) for input in inputs]
      xla_model((tuple(inputs_xla)))
      grads_output_xla = [
          torch_xla._XLAC.XLATensor(grad_output)
          for grad_output in grad_outputs[:gradient.f_real_outputs]
      ]
      xla_model.backward((tuple(grads_output_xla)))
      grad_inputs_xla = [input_xla.grad.to_tensor() for input_xla in inputs_xla]
      grad_inputs_xla.extend(
          [p.grad.to_tensor() for p in xla_model.parameters()[0]])
    ##############################################################
    # forward + backward with regular autograd / torch
    outputs_gt = model(*inputs)
    outputs_gt = xu.as_list(outputs_gt)
    grad_inputs_gt = torch.autograd.grad(
        outputs_gt, inputs_params, grad_outputs, only_inputs=True)
    for out_jit, out_autograd in zip(outputs, outputs_gt):
      self.assertEqualRel(
          out_jit, out_autograd, rel_err=rel_err, abs_err=abs_err)

    for grad_input_jit, grad_input_autograd in zip(grad_inputs, grad_inputs_gt):
      self.assertEqualRel(
          grad_input_jit, grad_input_autograd, rel_err=rel_err, abs_err=abs_err)

    # TODO: test buffers as well (running_mean, etc.)
    if xla:
      for i, (grad_input_jit, grad_input_xla) in enumerate(
          zip(grad_inputs, grad_inputs_xla)):
        self.assertEqualRel(grad_input_jit, grad_input_xla, rel_err, abs_err)

  def test_avgpool(self):

    class AvgPoolGrad(nn.Module):

      def __init__(self, stride, padding, count_include_pad):
        super(AvgPoolGrad, self).__init__()
        self.stride = stride
        self.padding = padding
        self.count_include_pad = count_include_pad

      def forward(self, x):
        return F.avg_pool2d(x, 2, self.stride, self.padding, False,
                            self.count_include_pad)

    for stride in [1, 2]:
      for padding in [0, 1]:
        for count_include_pad in [False, True]:
          model = AvgPoolGrad(stride, padding, count_include_pad)
          inputs = [_gen_tensor(4, 1, 28, 28, requires_grad=True)]
          self.checkGrad(model, inputs, xla=True)

  def test_adaptive_avgpool(self):

    class AdaptiveAvgPoolGrad(nn.Module):

      def __init__(self, output_size):
        super(AdaptiveAvgPoolGrad, self).__init__()
        self.output_size = output_size

      def forward(self, x):
        return F.adaptive_avg_pool2d(x, self.output_size)

    model = AdaptiveAvgPoolGrad((2, 3))
    for scale in [1, 2]:
      inputs = [_gen_tensor(10, 3, 2 * scale, 3 * scale, requires_grad=True)]
      self.checkGrad(model, inputs, xla=True)

  def test_threshold(self):

    class ThresholdPoolGrad(nn.Module):

      def __init__(self):
        super(ThresholdPoolGrad, self).__init__()
        self.threshold = nn.Threshold(0.4, 20)

      def forward(self, x):
        return self.threshold(x)

    model = ThresholdPoolGrad()
    inputs = [_gen_tensor(4, 2, requires_grad=True)]
    self.checkGrad(model, inputs, xla=True)

  def test_maxpool(self):

    class MaxPoolGrad(nn.Module):

      def forward(self, x):
        return F.max_pool2d(x, 2)

    model = MaxPoolGrad()
    inputs = [_gen_tensor(4, 1, 28, 28, requires_grad=True)]
    self.checkGrad(model, inputs, xla=True)

  def test_tanh(self):

    class TanhGrad(nn.Module):

      def forward(self, x):
        return torch.tanh(x)

    model = TanhGrad()
    inputs = [_gen_tensor(4, 2, requires_grad=True)]
    self.checkGrad(model, inputs, xla=True)

  def test_sigmoid(self):

    class SigmoidGrad(nn.Module):

      def forward(self, x):
        return torch.sigmoid(x)

    model = SigmoidGrad()
    inputs = [_gen_tensor(4, 2, requires_grad=True)]
    self.checkGrad(model, inputs, xla=True, rel_err=1e-2, abs_err=1e-2)

  @unittest.skip(
      'differentiation of prim::ListUnpack is not supported, or it is missing '
      'necessary type information')
  def test_chunk(self):

    class ChunkGrad(nn.Module):

      def forward(self, x):
        return x.chunk(2, 1)

    model = ChunkGrad()
    inputs = [_gen_tensor(4, 4, requires_grad=True)]
    self.checkGrad(model, inputs, xla=True)

  @unittest.skip('bool value of Tensor with more than one value is ambiguous')
  def test_lstm_cell(self):

    class LSTMCellGrad(nn.Module):

      def __init__(self):
        super(LSTMCellGrad, self).__init__()
        self.i2h = nn.Linear(3, 8)
        self.h2h = nn.Linear(2, 8)

      def forward(self, x, hx, cx):
        gates = self.i2h(x) + self.h2h(hx)

        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
        ingate = torch.sigmoid(ingate)
        forgetgate = torch.sigmoid(forgetgate)
        cellgate = torch.tanh(cellgate)
        outgate = torch.sigmoid(outgate)

        cy = (forgetgate * cx) + (ingate * cellgate)
        hy = outgate * torch.tanh(cy)
        return hy, cy

    model = LSTMCellGrad()
    inputs = [
        _gen_tensor(4, 3, requires_grad=True),
        _gen_tensor(4, 2, requires_grad=True),
        _gen_tensor(4, 2, requires_grad=True)
    ]
    self.checkGrad(model, inputs, xla=True)

  def test_conv2d(self):
    if FLAGS.long_test:
      config = [
          [1, 7, 15, 32],  # ichans
          [1, 4, 21, 32],  # ochans
          [1, 2, 3, 5],  # size
          [1, 2],  # stride
          [0, 1],  # padding
          [True, False],  # bias
      ]
    else:
      config = [
          [1, 5],  # ichans
          [1, 4],  # ochans
          [1, 3],  # size
          [1],  # stride
          [0],  # padding
          [False],  # bias
      ]
    for ichans, ochans, size, stride, padding, bias in (
        itertools.product(*config)):
      # TODO: dilation, groups, transpose
      model = nn.Conv2d(ichans, ochans, size, stride, padding, bias=bias)
      inputs = [_gen_tensor(4, ichans, 28, 28, requires_grad=True)]
      self.checkGrad(model, inputs, xla=True, abs_err=1e-3)

  def test_batchnorm2d(self):
    for chans in [1, 15, 32]:
      for eps in [1e-5, 1e-3, 1e-2]:
        # TODO: momentum, training, affine
        model = nn.BatchNorm2d(chans, eps=eps)
        inputs = [_gen_tensor(4, chans, 28, 28, requires_grad=True)]
        self.checkGrad(model, inputs, xla=True)

  def test_logsoftmax(self):
    for dim in [0, 1]:  # todo test 3d as well
      for batch in [1, 3, 4]:

        class LSMGrad(nn.Module):

          def forward(self, x):
            return F.log_softmax(x, dim)

        model = LSMGrad()
        inputs = [_gen_tensor(batch, 9, requires_grad=True)]
        self.checkGrad(model, inputs, xla=True)

  def test_nll_loss(self):
    input = _gen_tensor(3, 5, requires_grad=True)
    target = torch.empty(3, dtype=torch.long).random_(5)
    model = XlaNllLoss()
    traced_model = torch.jit.trace(model, (input, target))
    xla_model = torch_xla._XLAC.XlaModule(traced_model)
    xla_inputs = [
        torch_xla._XLAC.XLATensor(input),
        torch_xla._XLAC.XLATensor(target)
    ]
    output_xla = xla_model((tuple(xla_inputs)))
    xla_model.backward(*output_xla)
    output = model(input, target)
    output.backward()
    self.assertEqual(input.grad.data, xla_inputs[0].grad.data.to_tensor())

  def test_mnist(self):
    model = XlaMNIST()
    inputs = [_gen_tensor(4, 1, 28, 28, requires_grad=True)]
    self.checkGrad(model, inputs, xla=True)

  @unittest.skip('Disable until we figure out the precision issue')
  def test_resnet(self):
    model = torchvision.models.resnet18()
    inputs = [_gen_tensor(4, 3, 224, 224, requires_grad=True)]
    self.checkGrad(model, inputs, xla=False)

class TestOptimizer(XlaTestCase):

  def test_inplace_add_mul(self):
    orig_x = torch.Tensor([[1, 2, 3], [4, 5, 6]])
    orig_y = torch.Tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
    x = orig_x
    y = orig_y
    xla_x = torch_xla._XLAC.XLATensor(orig_x)
    xla_y = torch_xla._XLAC.XLATensor(orig_y)
    self.assertEqualDbg(
        x.add_(2, y).mul_(y),
        xla_x.add_(2, xla_y).mul_(xla_y).to_tensor())
    self.assertEqualDbg(
        x.add_(y).mul_(y),
        xla_x.add_(xla_y).mul_(xla_y).to_tensor())

  def test_add_mul(self):
    orig_x = torch.Tensor([[1, 2], [3, 4]])
    orig_y = torch.Tensor([[0.1, 0.2], [0.3, 0.4]])
    x = orig_x
    y = orig_y
    xla_x = torch_xla._XLAC.XLATensor(orig_x)
    xla_y = torch_xla._XLAC.XLATensor(orig_y)
    xla_ones = torch_xla._XLAC.XLATensor(torch.ones_like(x))
    self.assertEqualDbg(x + 3 * y, xla_x.add(3, xla_y).to_tensor())
    self.assertEqualDbg(x * y, xla_x.mul(xla_y).to_tensor())
    z = (x + 9) * (y + 3)
    xla_z = xla_x.add(9, xla_ones).mul(xla_y.add(3, xla_ones))
    self.assertEqualDbg(z, xla_z.to_tensor())
    self.assertEqualDbg(x + y, (xla_x + xla_y).to_tensor())
    self.assertEqualDbg(x * y, (xla_x * xla_y).to_tensor())
    self.assertEqualDbg(x * 11.0, (xla_x * 11.0).to_tensor())
    self.assertEqualDbg(x / 3.11, (xla_x / 3.11).to_tensor())
    self.assertEqualDbg(y / x, (xla_y / xla_x).to_tensor())

  def checkSgd(self, lr, momentum, weight_decay, nsteps, do_zero_grad):
    input = _gen_tensor(4, 4, requires_grad=True)
    model = nn.Linear(4, 20)
    traced_model = torch.jit.trace(model, input)
    xla_model = torch_xla._XLAC.XlaModule(traced_model)
    input_xla = [torch_xla._XLAC.XLATensor(input)]
    xla_model((tuple(input_xla)))
    xla_optimizer = optim.SGD(
        xla_model.parameters()[0],
        lr=lr,
        momentum=momentum,
        weight_decay=weight_decay)
    optimizer = optim.SGD(
        model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)
    output = model(input)
    grad_output = _gen_tensor(*output.shape)  # random gradients
    grad_output_xla = [torch_xla._XLAC.XLATensor(grad_output)]
    output.backward(grad_output)
    xla_model.backward((tuple(grad_output_xla)))
    if do_zero_grad:
      optimizer.zero_grad()
      xla_optimizer.zero_grad()
    for _ in range(0, nsteps):
      xla_optimizer.step()
      optimizer.step()
      xla_updated_params = [
          p.to_tensor().data for p in xla_model.parameters()[0]
      ]
      updated_params = [p.data for p in model.parameters()]
      for i in range(0, len(updated_params)):
        self.assertEqualRel(xla_updated_params[i], updated_params[i])

  def test_sgd(self):
    for weight_decay in [0, 5e-4]:
      self.checkSgd(
          lr=0.1,
          momentum=0,
          weight_decay=weight_decay,
          nsteps=1,
          do_zero_grad=True)
      self.checkSgd(
          lr=0.1,
          momentum=0,
          weight_decay=weight_decay,
          nsteps=2,
          do_zero_grad=False)
      self.checkSgd(
          lr=0.1,
          momentum=0.5,
          weight_decay=weight_decay,
          nsteps=1,
          do_zero_grad=True)
      self.checkSgd(
          lr=0.1,
          momentum=0.5,
          weight_decay=weight_decay,
          nsteps=2,
          do_zero_grad=False)

# Disabled always for now.
@unittest.skipIf(not (FLAGS.replicated and _support_replicated('TPU', 8)),
                 'Replicated (8) TPU only')
class TestReplicatedSum(XlaTestCase):

  def test(self):

    class XlaSum(nn.Module):

      def forward(self, x, y):
        return x + y

    model = XlaSum()
    for num_replicas in [2, 3, 4, 5, 6, 7, 8]:
      inputs = _random_inputs(((3, 3), (3, 3)), num_replicas=num_replicas)
      out = _xla_run(model, inputs)
      self.compareReplicated(model, inputs, out)

class TestSelect(XlaTestCase):

  def test_get_xla_tensor(self):
    x = _gen_tensor(14, 24, 8, device=xm.xla_device())
    t = x.data.cpu()
    sx = x.select(1, 12)
    tx = t.select(1, 12)
    self.assertEqual(tx, sx.data.cpu())

class TestXLATensor(XlaTestCase):

  def test_size(self):
    x = _gen_tensor(2, 1, 4, 6)
    xt_x = torch_xla._XLAC.XLATensor(x)
    rank = x.dim()
    for dim in range(-rank, rank):
      self.assertEqual(x.size(dim), xt_x.size(dim))

  def test_relu(self):
    x = _gen_tensor(2, 1, 4, 6)
    xt_x = torch_xla._XLAC.XLATensor(x)
    expected = F.relu(x)
    out = torch_xla._XLAC.relu(xt_x).to_tensor()
    self.assertEqualDbg(out.data, expected.data)

  def test_threshold(self):
    x = _gen_tensor(2, 1, 4, 6)
    xt_x = torch_xla._XLAC.XLATensor(x)
    threshold = 0.4
    value = 20
    expected = F.threshold(x, threshold, value)
    out = torch_xla._XLAC.threshold(xt_x, threshold, value).to_tensor()
    self.assertEqualDbg(out.data, expected.data)

  def test_conv2d(self):
    in_channels = 3
    out_channels = 7
    kernel_size = 5
    input = _gen_tensor(4, in_channels, 28, 28)
    weight = _gen_tensor(out_channels, in_channels, kernel_size, kernel_size)
    bias = _gen_tensor(out_channels)
    xt_input = torch_xla._XLAC.XLATensor(input)
    xt_weight = torch_xla._XLAC.XLATensor(weight)
    xt_bias = torch_xla._XLAC.XLATensor(bias)
    for stride in range(1, 4):
      for padding in range(0, 3):
        for with_bias in [True, False]:
          conv_bias = bias if with_bias else None
          expected = F.conv2d(
              input, weight, conv_bias, stride=stride, padding=padding)
          if with_bias:
            out = torch_xla._XLAC.conv2d(
                xt_input, xt_weight, xt_bias, stride=stride,
                padding=padding).to_tensor()
          else:
            out = torch_xla._XLAC.conv2d(
                xt_input, xt_weight, stride=stride,
                padding=padding).to_tensor()
          self.assertEqualRel(out.data, expected.data)

  def test_conv2d_non_square(self):
    in_channels = 3
    out_channels = 7
    kernel_size = 5
    input = _gen_tensor(4, in_channels, 28, 28)
    weight = _gen_tensor(out_channels, in_channels, kernel_size, kernel_size)
    bias = _gen_tensor(out_channels)
    xt_input = torch_xla._XLAC.XLATensor(input)
    xt_weight = torch_xla._XLAC.XLATensor(weight)
    xt_bias = torch_xla._XLAC.XLATensor(bias)
    for stride in range(1, 4):
      for padding in range(0, 3):
        for with_bias in [True, False]:
          conv_bias = bias if with_bias else None
          expected = F.conv2d(
              input,
              weight,
              conv_bias,
              stride=[stride, stride + 1],
              padding=[padding, padding + 1])
          if with_bias:
            out = torch_xla._XLAC.conv2d(
                xt_input,
                xt_weight,
                xt_bias,
                stride=[stride, stride + 1],
                padding=[padding, padding + 1]).to_tensor()
          else:
            out = torch_xla._XLAC.conv2d(
                xt_input,
                xt_weight,
                stride=[stride, stride + 1],
                padding=[padding, padding + 1]).to_tensor()
          self.assertEqualRel(out.data, expected.data)

  def test_addmm(self):
    in_channels = 32
    out_channels = 320
    labels = 50
    input = _gen_tensor(in_channels, out_channels)
    weight = _gen_tensor(out_channels, labels)
    bias = _gen_tensor(labels)
    xt_input = torch_xla._XLAC.XLATensor(input)
    xt_weight = torch_xla._XLAC.XLATensor(weight)
    xt_bias = torch_xla._XLAC.XLATensor(bias)
    out = torch.addmm(bias, input, weight)
    expected = torch_xla._XLAC.addmm(xt_bias, xt_input, xt_weight).to_tensor()
    self.assertEqualRel(out.data, expected.data)

  def test_max_pool2d(self):
    x = _gen_tensor(1, 64, 112, 112)
    xt_x = torch_xla._XLAC.XLATensor(x)
    for stride in [1, 2]:
      for padding in [0, 1]:
        expected = F.max_pool2d(x, 3, stride=stride, padding=padding)
        out = torch_xla._XLAC.max_pool2d(
            xt_x, 3, stride=stride, padding=padding).to_tensor()
        self.assertEqualRel(out.data, expected.data)

  def test_max_pool2d_non_square(self):
    x = _gen_tensor(1, 64, 112, 112)
    xt_x = torch_xla._XLAC.XLATensor(x)
    for stride in [1, 2]:
      for padding in [0, 1]:
        expected = F.max_pool2d(
            x, [3, 4],
            stride=[stride, stride + 1],
            padding=[padding, padding + 1])
        out = torch_xla._XLAC.max_pool2d(
            xt_x, [3, 4],
            stride=[stride, stride + 1],
            padding=[padding, padding + 1]).to_tensor()
        self.assertEqualRel(out.data, expected.data)

  def test_avg_pool2d(self):
    x = _gen_tensor(4, 1, 28, 28)
    xt_x = torch_xla._XLAC.XLATensor(x)
    for stride in [1, 2]:
      for padding in [0, 1]:
        for count_include_pad in [False, True]:
          expected = F.avg_pool2d(
              x,
              2,
              stride=stride,
              padding=padding,
              count_include_pad=count_include_pad)
          out = torch_xla._XLAC.avg_pool2d(
              xt_x,
              2,
              stride=stride,
              padding=padding,
              count_include_pad=count_include_pad).to_tensor()
          self.assertEqualRel(out.data, expected.data)

  def test_avg_pool2d_non_square(self):
    x = _gen_tensor(4, 1, 28, 28)
    xt_x = torch_xla._XLAC.XLATensor(x)
    for stride in [1, 2]:
      for padding in [0, 1]:
        for count_include_pad in [False, True]:
          expected = F.avg_pool2d(
              x, [4, 5],
              stride=[stride, stride + 1],
              padding=[padding, padding + 1],
              count_include_pad=count_include_pad)
          out = torch_xla._XLAC.avg_pool2d(
              xt_x, [4, 5],
              stride=[stride, stride + 1],
              padding=[padding, padding + 1],
              count_include_pad=count_include_pad).to_tensor()
          self.assertEqualRel(out.data, expected.data)

  def test_transpose(self):
    x = _gen_tensor(2, 3)
    xt_x = torch_xla._XLAC.XLATensor(x)
    expected = x.t()
    out = xt_x.t().to_tensor()
    self.assertEqualDbg(out.data, expected.data)

  def test_view(self):
    x = _gen_tensor(32, 20, 4, 4)
    xt_x = torch_xla._XLAC.XLATensor(x)
    expected = x.view(-1, 320)
    out = xt_x.view(-1, 320).to_tensor()
    self.assertEqualDbg(out.data, expected.data)

  def test_log_softmax(self):
    x = _gen_tensor(5, 3, 4, 2)
    xt_x = torch_xla._XLAC.XLATensor(x)
    for dim in range(0, x.dim()):
      expected = x.log_softmax(dim)
      out = xt_x.log_softmax(dim).to_tensor()
      self.assertEqualRel(out.data, expected.data)

class TestAtenXlaTensor(XlaTestCase):

  def test_get_xla_tensor(self):
    t = _gen_tensor(4, 2, device=xm.xla_device())
    x = torch_xla._XLAC._get_xla_tensor(t)
    self.assertEqual(t.data.cpu(), x.to_tensor())

  def test_negative_slice(self):
    t = _gen_tensor(32, 24, 32)
    x = t.to(xm.xla_device())
    t_slice = t[:, :, -1]
    x_slice = x[:, :, -1]
    self.assertEqual(t_slice.data, x_slice.data.cpu())

  def test_masked_fill_with_tensor(self):
    input = _gen_tensor(2, 5, 4, 3)
    mask = torch.randint(0, 2, input.size(), dtype=torch.uint8)
    value = torch.tensor(42)
    xla_input = input.to(xm.xla_device())
    xla_mask = mask.to(xm.xla_device())
    xla_value = value.to(xm.xla_device())
    result = torch.masked_fill(input, mask, value)
    xla_result = torch.masked_fill(xla_input, xla_mask, xla_value)
    self.assertEqual(input.data, xla_input.data.cpu())
    self.assertEqual(result.data, xla_result.data.cpu())

  def test_add_mixed_device(self):
    input = _gen_tensor(3, 800, 1066)
    xla_input = input.to(xm.xla_device())
    output = input + 2
    xla_output = xla_input + 2
    self.assertEqual(output.data, xla_output.data.cpu())

  def test_mul_mixed_device(self):
    input = _gen_tensor(3, 800, 1066)
    xla_input = input.to(xm.xla_device())
    output = input * 2
    xla_output = xla_input * 2
    self.assertEqual(output.data, xla_output.data.cpu())

  def test_sub_mixed_device(self):
    input = _gen_tensor(3, 800, 1066)
    xla_input = input.to(xm.xla_device())
    output = input - 2
    xla_output = xla_input - 2
    self.assertEqual(output.data, xla_output.data.cpu())

  def test_div_mixed_device(self):
    input = _gen_tensor(3, 800, 1066)
    xla_input = input.to(xm.xla_device())
    output = input / 2
    xla_output = xla_input / 2
    self.assertEqual(output.data, xla_output.data.cpu())

if __name__ == '__main__':
  torch.set_default_tensor_type('torch.FloatTensor')
  torch.manual_seed(42)
  torch_xla._XLAC._xla_set_use_full_mat_mul_precision(
      use_full_mat_mul_precision=True)
  run_tests()


from django.http import JsonResponse
from models import *
from django.views.decorators.csrf import csrf_exempt
import json

# Create your views here.
def list_users(request):
    json_list = []
    users_list = Users.objects.all()
    for user in users_list:
        followers_raw = user.followed.all()
        followers = []
        for follower in followers_raw:
            followers.append(
                follower.firstName + ' ' + follower.lastName
            )
        json_list.append({
            'id': user.id,
            'firstName': user.firstName,
            'lastName': user.lastName,
            'followers': followers
        })
    return JsonResponse(json_list, safe=False)

@csrf_exempt
def follow(request):
    try:
        body = json.loads(request.body.decode('utf-8'))
        follower = Users.objects.get(id=body['follower'])
        followee = Users.objects.get(id=body['followee'])
        Follows.objects.create(
            follower=follower,
            followee=followee
        )
        return JsonResponse({
            'response': 'Successfully created follow'
        })
    except:
        return JsonResponse({
            'response': 'Something went wrong'
        })


#!/usr/bin/env python

# Define a schema using the DDL

from pyspark.sql import SparkSession
import pyspark.sql.types as st

# define schema for our data
schema = st.StructType(
    [
        st.StructField("Id", st.IntegerType(), False),
        st.StructField("First", st.StringType(), False),
        st.StructField("Last", st.StringType(), False),
        st.StructField("Url", st.StringType(), False),
        st.StructField("Published", st.StringType(), False),
        st.StructField("Hits", st.IntegerType(), False),
        st.StructField("Campaigns", st.ArrayType(st.StringType()), False),
    ]
)

# create our data
data = [
    [1, "Jules", "Damji", "https://tinyurl.1", "1/4/2016", 4535, ["twitter", "LinkedIn"], ],
    [2, "Brooke", "Wenig", "https://tinyurl.2", "5/5/2018", 8908, ["twitter", "LinkedIn"], ],
    [3, "Denny", "Lee", "https://tinyurl.3", "6/7/2019", 7659, ["web", "twitter", "FB", "LinkedIn"], ],
    [4, "Tathagata", "Das", "https://tinyurl.4", "5/12/2018", 10568, ["twitter", "FB"], ],
    [5, "Matei", "Zaharia", "https://tinyurl.5", "5/14/2014", 40578, ["web", "twitter", "FB", "LinkedIn"], ],
    [6, "Reynold", "Xin", "https://tinyurl.6", "3/2/2015", 25568, ["twitter", "LinkedIn"], ],
]

if __name__ == '__main__':
    spark = SparkSession.builder.appName('Example-3_6').getOrCreate()

    # create a DF using the above schema
    blogs_df = spark.createDataFrame(data, schema)

    # show the DF
    blogs_df.show()

    # print the schema used by Spark
    print('schema')
    print(blogs_df.printSchema())


import torch
from network import Net
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np

from data import send_to_device, to_torch_tensors, show_data

import os
import time
from tqdm import tqdm
import pdb
import utils
import pytorch_ssim

class Trainer():
    def __init__(self, args, train_loader, test_loader, writer = None):
        self.args = args
        self.model = args.model
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.device = args.device
        self.recon_kernel_size = self.args.recon_kernel_size
        self.eps = 0.00316
        self.global_step = 0
        self.model_dir = os.path.join('model', self.args.model)
        self.print_freq = self.args.print_freq
        self.writer = writer

        if not os.path.exists(self.model_dir): os.makedirs(self.model_dir)

        print('Making the Network')
        self.diffuseNetList = [Net(self.args).to(self.device)]*5
        # self.diffuseNet = Net(self.args).to(self.device)
        self.specularNet = Net(self.args).to(self.device)

    def train(self, epochs=200, learning_rate=1e-4, show_images=False):

        dataloader = self.train_loader

        # if self.args.debug:
        #     print(self.diffuseNet, "CUDA:", next(self.diffuseNet.parameters()).is_cuda)
        #     print(self.specularNet, "CUDA:", next(self.specularNet.parameters()).is_cuda)

        criterion = nn.L1Loss()
        # criterion = pytorch_ssim.SSIM()

        optimizerDiffList = []
        for i in range(len(self.diffuseNetList)):
            optimizerDiffList.append(optim.Adam(self.diffuseNetList[i].parameters(), lr=learning_rate))
        # optimizerDiff = optim.Adam(self.diffuseNet.parameters(), lr=learning_rate)
        optimizerSpec = optim.Adam(self.specularNet.parameters(), lr=learning_rate)
        schedulerDiffList = []
        for i in range(len(optimizerDiffList)):
            schedulerDiffList.append(torch.optim.lr_scheduler.MultiStepLR(schedulerDiffList[i], [10**4], gamma=0.01))
        schedulerSpec = torch.optim.lr_scheduler.MultiStepLR(optimizerSpec, [10**4], gamma=0.01)
        # schedulerDiff = torch.optim.lr_scheduler.StepLR(optimizerDiff, step_size=10**4, gamma=0.1)
        # schedulerSpec = torch.optim.lr_scheduler.StepLR(optimizerSpec, step_size=10**4, gamma=0.1)

        if self.args.resume:
            for i in range(len(self.diffuseNetList)):
                self.diffuseNetList[i].load_state_dict(torch.load(os.path.join(self.model_dir, 'diffuseNet{}.pt'.format(i))))
            self.specularNet.load_state_dict(torch.load(os.path.join(self.model_dir, 'specularNet.pt')))

        writer_LossDiff = 0
        writer_LossSpec = 0
        writer_LossFinal = 0
        writer_SSIMDiff = 0
        writer_SSIMSpec = 0
        writer_SSIMFinal = 0
        tr_fin_best = 100

        # lDiff = []
        # lSpec = []
        # lFinal = []

        print('Start Training')
        start = time.time()
        permutation = [0, 3, 1, 2]

        # loader_start = time.time()
        for epoch in range(epochs):
            print('Epoch {:04d}'.format(epoch))
            for i_batch, sample_batched in enumerate(tqdm(dataloader, ncols = 80)):
                self.global_step += 1

                # loader_end = time.time()

                # get the inputs
                X_diff = sample_batched['X_diff'].permute(permutation).to(self.device)

                Y_diff = sample_batched['Reference'][:,:,:,:3].permute(permutation).to(self.device)

                # zero the parameter gradients
                optimizerDiff.zero_grad()

                # forward + backward + optimize
                outputDiff = self.diffuseNet(X_diff)
                # if torch.isinf(outputDiff).any():
                #     print("Found Infinity Values in Kernel, Going to Skip this Batch!")
                #     continue
                # if torch.isnan(outputDiff).any():
                #     print("Found NaN Values in Kernel, Going to Skip this Batch!")
                #     continue

                # print(outputDiff.shape)

                if self.model == 'KPCN':
                    X_input = self.crop_like(X_diff, outputDiff)
                    outputDiff = self.apply_kernel(outputDiff, X_input)

                # if torch.isinf(outputDiff).any():
                #     print("Found Infinity Values in Output, Going to Skip this Batch!")
                #     continue
                # if torch.isnan(outputDiff).any():
                #     print("Found NaN Values in Output, Going to Skip this Batch!")
                #     continue

                Y_diff = self.crop_like(Y_diff, outputDiff)

                # lossDiff = criterion(outputDiff * 255, Y_diff * 255)
                # lossDiff = -criterion(outputDiff, Y_diff)
                lossDiff = criterion(outputDiff, Y_diff)
                ssimDiff = pytorch_ssim.ssim(outputDiff, Y_diff)
                lossDiff.backward()
                 # torch.nn.utils.clip_grad_norm_(self.diffuseNet.parameters(), 5)
                # torch.nn.utils.clip_grad_value_(self.diffuseNet.parameters(), 10)

                # Check if gradient became NaN --> Skip
                # found_NaN = False
                # for param in self.diffuseNet.parameters():
                #     if not torch.isnan(param.grad.data).any():
                #         continue
                #     found_NaN = True
                #     break
                # if found_NaN:
                #     print("Found NaN Values in Gradient, Goint to Skip this Batch!")
                #     continue

                optimizerDiff.step()

                # get the inputs
                X_spec = sample_batched['X_spec'].permute(permutation).to(self.device)
                Y_spec = sample_batched['Reference'][:,:,:,3:6].permute(permutation).to(self.device)

                # zero the parameter gradients
                optimizerSpec.zero_grad()

                # forward + backward + optimize
                outputSpec = self.specularNet(X_spec)

                if self.model == 'KPCN':
                    X_input = self.crop_like(X_spec, outputSpec)
                    outputSpec = self.apply_kernel(outputSpec, X_input)

                Y_spec = self.crop_like(Y_spec, outputSpec)

                # lossSpec = criterion(outputSpec * 255, Y_spec * 255)
                # lossSpec = -criterion(outputSpec, Y_spec)
                lossSpec = criterion(outputSpec, Y_spec)
                ssimSpec = pytorch_ssim.ssim(outputSpec, Y_spec)
                lossSpec.backward()
                optimizerSpec.step()

                schedulerDiff.step()
                schedulerSpec.step()

                # calculate final ground truth error
                with torch.no_grad():
                    albedo = sample_batched['origAlbedo'].permute(permutation).to(self.device)
                    albedo = self.crop_like(albedo, outputDiff)
                    outputFinal = outputDiff * (albedo + self.eps) + torch.exp(outputSpec) - 1.0

                    # if False:#i_batch % 500:
                    #     print("Sample, denoised, gt")
                    #     sz = 3
                    #     orig = self.crop_like(sample_batched['finalInput'].permute(permutation), outputFinal)
                    #     orig = orig.cpu().permute([0, 2, 3, 1]).numpy()[0,:]
                    #     img = outputFinal.cpu().permute([0, 2, 3, 1]).numpy()[0,:]
                    #     gt = self.crop_like(sample_batched['finalGt'].permute(permutation), outputFinal)
                    #     gt = gt.cpu().permute([0, 2, 3, 1]).numpy()[0,:]

                    Y_final = sample_batched['finalGt'].permute(permutation).to(self.device)
                    Y_final = self.crop_like(Y_final, outputFinal)
                    lossFinal = criterion(outputFinal, Y_final)
                    ssimFinal = pytorch_ssim.ssim(outputFinal, Y_final)
                    # accuLossFinal += lossFinal.item()

                iter_lossDiff, iter_lossSpec, iter_lossFinal = lossDiff.item(), lossSpec.item(), lossFinal.item()
                # accuLossDiff += iter_lossDiff
                # accuLossSpec += iter_lossSpec
                writer_LossDiff += iter_lossDiff
                writer_LossSpec += iter_lossSpec
                writer_LossFinal += iter_lossFinal
                writer_SSIMDiff += 1-ssimDiff.item()
                writer_SSIMSpec += 1-ssimSpec.item()
                writer_SSIMFinal += 1-ssimFinal.item()
                # loader_start = time.time()

                if (self.global_step-1) % self.print_freq == 0:
                    self.writer.add_scalars('data/L1Loss', {'LossDiff': writer_LossDiff/self.print_freq,
                                                        'LossSpec':writer_LossSpec/self.print_freq,
                                                        'LossFinal':writer_LossFinal/self.print_freq}, self.global_step)
                    self.writer.add_scalars('data/SSIM', {'SSIMDiff': writer_SSIMDiff/self.print_freq,
                                                        'SSIMSpec':writer_SSIMSpec/self.print_freq,
                                                        'SSIMFinal':writer_SSIMFinal/self.print_freq}, self.global_step)
                    # self.writer.add_scalars('data/LossSpec', {self.model: writer_LossSpec/self.print_freq}, self.global_step)
                    # self.writer.add_scalars('data/LossFinal', {self.model: writer_LossFinal/self.print_freq}, self.global_step)
                    # self.writer.add_scalars('data/SSIMDiff', {self.model: writer_SSIMDiff/self.print_freq}, self.global_step)
                    # self.writer.add_scalars('data/SSIMSpec', {self.model: writer_SSIMSpec/self.print_freq}, self.global_step)
                    # self.writer.add_scalars('data/SSIMFinal', {self.model: writer_SSIMFinal/self.print_freq}, self.global_step)
                    writer_LossDiff, writer_LossSpec, writer_LossFinal = 0, 0, 0
                    writer_SSIMDiff, writer_SSIMSpec, writer_SSIMFinal = 0, 0, 0

                    if tr_fin_best > writer_LossFinal/self.print_freq:
                        tr_fin_best = writer_LossFinal/self.print_freq
                        torch.save(self.diffuseNet.state_dict(), os.path.join(self.model_dir, 'fin_best_diff.pt'))
                        torch.save(self.specularNet.state_dict(), os.path.join(self.model_dir, 'fin_best_spec.pt'))

                if self.global_step % self.args.save_freq == 0:
                    torch.save(self.diffuseNet.state_dict(), os.path.join(self.model_dir, 'diffuseNet.pt'))
                    torch.save(self.specularNet.state_dict(), os.path.join(self.model_dir, 'specularNet.pt'))

                if self.global_step % self.args.test_freq == 0:
                    self.test()

            # print("Epoch {}".format(epoch + 1))
            # print("LossDiff: {}".format(accuLossDiff))
            # print("LossSpec: {}".format(accuLossSpec))
            # print("LossFinal: {}".format(accuLossFinal))

            # lDiff.append(accuLossDiff)
            # lSpec.append(accuLossSpec)
            # lFinal.append(accuLossFinal)

            # accuLossDiff = 0
            # accuLossSpec = 0
            # accuLossFinal = 0

        print('Finished training in mode', self.model)
        print('Took', time.time() - start, 'seconds.')

        # return self.diffuseNet, self.specularNet, lDiff, lSpec, lFinal

    def crop_like(self, data, like, debug=False):
        if data.shape[-2:] != like.shape[-2:]:
            # crop
            with torch.no_grad():
                dx, dy = data.shape[-2] - like.shape[-2], data.shape[-1] - like.shape[-1]
                data = data[:,:,dx//2:-dx//2,dy//2:-dy//2]
                if debug:
                    print(dx, dy)
                    print("After crop:", data.shape)
        return data

    def apply_kernel(self, weights, data):
        # apply softmax to kernel weights
        weights = weights.permute((0, 2, 3, 1)).to(self.device)
        _, _, h, w = data.size()
        weights = F.softmax(weights, dim=3).view(-1, w * h, self.recon_kernel_size, self.recon_kernel_size)

        # now we have to apply kernels to every pixel
        # first pad the input
        r = self.recon_kernel_size // 2
        data = F.pad(data[:,:3,:,:], (r,) * 4, "reflect")

        #print(data[0,:,:,:])

        # make slices
        R = []
        G = []
        B = []
        kernels = []
        for i in range(h):
            for j in range(w):
                pos = i*h+j
                ws = weights[:,pos:pos+1,:,:]
                kernels += [ws, ws, ws]
                sy, ey = i+r-r, i+r+r+1
                sx, ex = j+r-r, j+r+r+1
                R.append(data[:,0:1,sy:ey,sx:ex])
                G.append(data[:,1:2,sy:ey,sx:ex])
                B.append(data[:,2:3,sy:ey,sx:ex])
                #slices.append(data[:,:,sy:ey,sx:ex])

        reds = (torch.cat(R, dim=1).to(self.device)*weights).sum(2).sum(2)
        greens = (torch.cat(G, dim=1).to(self.device)*weights).sum(2).sum(2)
        blues = (torch.cat(B, dim=1).to(self.device)*weights).sum(2).sum(2)

        res = torch.cat((reds, greens, blues), dim=1).view(-1, 3, h, w).to(self.device)

        return res

    def unsqueeze_all(self, d):
        for k, v in d.items():
            d[k] = torch.unsqueeze(v, dim=0)
        return d

    def test(self):
        if self.args.only_test:
            self.diffuseNet.load_state_dict(torch.load(os.path.join(self.model_dir, 'diffuseNet.pt')))
            self.specularNet.load_state_dict(torch.load(os.path.join(self.model_dir, 'specularNet.pt')))
        for i_batch, tp in enumerate(tqdm(self.test_loader)):
            img_num, test_data = tp
            self.denoise(test_data, img_num[0])

    def denoise(self, data, img_num, debug=False):
        permutation = [0, 3, 1, 2]
        eps = 0.00316
        diffuseNet = self.diffuseNet
        specularNet = self.specularNet
        with torch.no_grad():
            # pdb.set_trace()
            out_channels = diffuseNet.net[len(diffuseNet.net)-1].out_channels
            mode = 'DPCN' if out_channels == 3 else 'KPCN'
            criterion = nn.L1Loss()
            # criterion = pytorch_ssim.SSIM()

            if debug:
                print("Out channels:", out_channels)
                print("Detected mode", mode)

            # make singleton batch
            data = send_to_device(to_torch_tensors(data), self.device)
            if len(data['X_diff'].size()) != 4:
                data = self.unsqueeze_all(data)
            if debug:
                print(data['X_diff'].size())

            X_diff = data['X_diff'].permute(permutation).to(self.device)
            Y_diff = data['Reference'][:,:,:,:3].permute(permutation).to(self.device)

            # forward + backward + optimize
            outputDiff = diffuseNet(X_diff)

            # print(outputDiff.shape)

            if mode == 'KPCN':
                X_input = self.crop_like(X_diff, outputDiff)
                outputDiff = self.apply_kernel(outputDiff, X_input)

            Y_diff = self.crop_like(Y_diff, outputDiff)

            lossDiff = criterion(outputDiff, Y_diff).item()
            ssimDiff = 1-pytorch_ssim.ssim(outputDiff, Y_diff)

            # get the inputs
            X_spec = data['X_spec'].permute(permutation).to(self.device)
            Y_spec = data['Reference'][:,:,:,3:6].permute(permutation).to(self.device)

            # forward + backward + optimize
            outputSpec = specularNet(X_spec)

            if mode == 'KPCN':
                X_input = self.crop_like(X_spec, outputSpec)
                outputSpec = self.apply_kernel(outputSpec, X_input)

            Y_spec = self.crop_like(Y_spec, outputSpec)

            lossSpec = criterion(outputSpec, Y_spec).item()
            ssimSpec = 1-pytorch_ssim.ssim(outputSpec, Y_spec)

            # calculate final ground truth error
            albedo = data['origAlbedo'].permute(permutation).to(self.device)
            albedo = self.crop_like(albedo, outputDiff)
            outputFinal = outputDiff * (albedo + eps) + torch.exp(outputSpec) - 1.0

            # save image
            img = outputFinal.cpu().permute([0, 2, 3, 1]).numpy()[0,:]
            img = np.clip(img, 0, 1)**0.45454545
            new_img = (img * 255 / np.max(img)).astype('uint8')
            new_img1 = np.transpose(new_img, [2, 0, 1])

            Y_final = data['finalGt'].permute(permutation).to(self.device)
            Y_final = self.crop_like(Y_final, outputFinal)

            # save image
            img = Y_final.cpu().permute([0, 2, 3, 1]).numpy()[0,:]
            img = np.clip(img, 0, 1)**0.45454545
            new_img = (img * 255 / np.max(img)).astype('uint8')
            new_img2 = np.transpose(new_img, [2, 0, 1])

            Y_orig = data['finalInput'].permute(permutation).to(self.device)
            Y_orig = self.crop_like(Y_orig, outputFinal)
            img = Y_orig.cpu().permute([0, 2, 3, 1]).numpy()[0,:]
            img = np.clip(img, 0, 1)**0.45454545
            new_img = (img * 255 / np.max(img)).astype('uint8')
            new_img3 = np.transpose(new_img, [2, 0, 1])

            new_img = np.concatenate((new_img3, new_img1, new_img2), axis=2)
            self.writer.add_image('img/'+str(img_num), new_img, self.global_step)

            lossFinal = criterion(outputFinal, Y_final).item()
            ssimFinal = 1-pytorch_ssim.ssim(outputFinal, Y_final)
            # psnrFinal = utils.calc_psnr(outputFinal, Y_final)

            self.writer.add_scalars('data/l1', {'{}_lossDiff'.format(img_num): lossDiff, '{}_lossSpec'.format(img_num) : lossSpec, '{}_lossFinal'.format(img_num) : lossFinal}, self.global_step)
            self.writer.add_scalars('data/ssim', {'{}_ssimDiff'.format(img_num): ssimDiff, '{}_ssimSpec'.format(img_num):ssimSpec, '{}_ssimFinal'.format(img_num):ssimFinal}, self.global_step)


import json
from collections import defaultdict

def create_dict_with_all_stops(response):
    """
    Funkcja wyciąga z responsu wszystkie przystanki i tworzy słownik.
    :param response: wynik zapytania API o listę przystanków.
    :return: słownik ze wszystkimi przystankami, w którym kluczem jest nazwa, a wartością id.
    """
    all_stops_dict = defaultdict(str)
    json_format = json.loads(response)
    json_values = json_format["stops"]
    for x in json_values:
        all_stops_dict[x["name"]] = x["shortName"]

    return all_stops_dict

def create_dict_with_results(response, user_input):
    """
    Funkcja filtruje response zawierający wszystkie przyjazdy na podstawie user_input
    i zwraca słownik z wynikami filtrowania.
    :param response: wynik zapytania API o przyjazdy na przystanek.
    :param user_input: słownik z przystankiem, numerem linii i kierunkiem.
    :return: zwraca słownik zawierający słowniki z poszczególnymi wynikami filtrowania.
    """
    results = defaultdict(dict)
    json_format = json.loads(response)
    json_values = json_format["actual"]
    json_values = filter_results(json_values, user_input)
    index = 1
    for x in json_values:
        results[index] = {"LineNumber": x["patternText"], "Direction": x["direction"], "Time": x["actualRelativeTime"]}
        index += 1

    return results

def filter_results(json_values, user_input):
    """
    Funkcja filtruje przyjazdy na dany przystanek po numerze linii i kierunku, i zwraca liste z wynikami filtrowania.
    :param json_values: lista słowników z wynikami zapytania o przyjazdy tramwajów na przystanek.
    :param user_input: słownik z przystankiem, numerem linii i kierunkiem.
    :return: zwraca listę ze słownikami spełniającymi kryteria.
    """
    list_with_filtered_results = []
    for value in json_values:
        if user_input["LineNumber"]:
            if user_input["LineNumber"] != value["patternText"]:
                continue

        if user_input["Direction"]:
            if user_input["Direction"] != value["direction"]:
                continue

        if value["actualRelativeTime"] > 0:
            list_with_filtered_results.append(value)

    return list_with_filtered_results


# -*- coding: utf-8 -*-
import time
import requests
from Utility.Baidu_Message.log import log
from Utility.Baidu_Message.config import TOKEN_GET_API, TOKEN_TEST_API, TOKEN_CLIENT_ID, TOKEN_CLIENT_SECRET, \
    ADSL_APIURL, ADSL_USERNAME, ADSL_PASSWORD

class TokenGen:
    def __init__(self, client_id, client_secret, **kwargs):
        """
        :param client_id: 传入的个人 `client_id`
        :param client_secret: 传入的个人 `client_secret`
        :param kwargs:
        """
        self.flag = True
        self.timing = time.time()
        self.client_id = client_id
        self.client_secret = client_secret
        self.token_data = {'client_id': client_id,
                           'client_secret': client_secret,
                           'grant_type': 'client_credentials',
                           'scope': 'getproxyip',
                           }
        # token获取API
        self.token_get_api = TOKEN_GET_API
        # token验证API
        self.token_test_api = TOKEN_TEST_API
        self.token = ''

    def get_token(self):
        """
        当flag为True时获取token
        当flag为False时 - token时效 > 1 hour重新获取新token
                       - token时效 < 1 hour直接返回旧token
        :return:
        """
        if not self.token or self.flag or time.time() - self.timing > 60 * 60:
            self.token = self._request_token()
            # 更新获取到的新token的时间
            self.timing = time.time()
        # 测试token是否可用
        self._test_token(new_token=self.token)
        if self.flag:
            return self.get_token()
        return self.token

    def _request_token(self):
        """
        请求获取token
        :return:
        """
        headers = {'Connection': 'close', }
        try:
            resp = requests.post(url=self.token_get_api, data=self.token_data, headers=headers, timeout=5)
            if resp.status_code == 200:
                resp_json = resp.json()
                return resp_json.get('token_type') + ' ' + resp_json.get('access_token')
        except requests.exceptions.RequestException:
            pass
        return None

    def _test_token(self, new_token):
        """
        测试token
        :param new_token:
        :return:
        """
        headers = {
            'Connection': 'close',
            'Authorization': new_token,
        }
        try:
            resp = requests.get(url=self.token_test_api, headers=headers, timeout=5)
            # 如果状态码401则token失效
            if resp.status_code == 401:
                self.flag = True
            # token可用
            else:
                self.flag = False
        # 验证过程中异常则视为不可用
        except Exception as e:
            log.error('测试token失败...{}'.format(e))
            self.flag = True

def get_ADSL_proxy():
    # 不使用代理 直接放回 {}
    # return {}
    # time.sleep(5)
    token_gen = TokenGen(client_id=TOKEN_CLIENT_ID, client_secret=TOKEN_CLIENT_SECRET)
    proxies = {}
    try:
        headers = {
            'Connection': 'close',
            'Authorization': token_gen.get_token(),
        }
        res = requests.get(url=ADSL_APIURL, headers=headers, timeout=15)
        ip = res.json().get('IP')
        port = res.json().get('Port')
        if ip and port:
            proxy = 'http://{}:{}@{}:{}'.format(ADSL_USERNAME, ADSL_PASSWORD, ip, port)
            proxies = {
                "http": proxy,
                "https": proxy,
            }
            # log.info('获取代理成功...{}'.format(proxies))
    except Exception as e:
        log.error('获取代理失败...{}'.format(e))
    return proxies

if __name__ == '__main__':
    # for i in range(3):
    #     temp_token = token_gen.get_token()
    #     log.info(temp_token)
    #     time.sleep(2)
    info_log.info(get_ADSL_proxy())
    url = 'https://www.toutiao.com/a6639533048863916552'
    # for i in range(100):
    #
    #     r = requests.get(url)
    #     print(r.text)
    get_adsl = get_ADSL_proxy()
    print(get_adsl)
    # r = requests.get(url, proxies=proxies)
    # print(r.text)


from requests.exceptions import HTTPError

from src.core import api
from src.core.filters import create_datetime

try:
    # Get the date of the prompt
    prompt_date = input("Enter the prompt date (YYYY-MM-DD): ")

    # Send out a broadcast
    api.post("subscription", "broadcast", params={"date": create_datetime(prompt_date)})
    print(f"Email broadcast for {prompt_date} successfully sent")

# A broadcast for that day couldn't be sent
except HTTPError:
    print(f"Unable to send email broadcast for {prompt_date}!")
raise SystemExit(0)


from __future__ import division,print_function,absolute_import
import tflearn
from tflearn.layers.core import *
from tflearn.layers.conv import *
from tflearn.layers.estimator import *
from tflearn.data_utils import *

# Load Cifar10 DataSet
from tflearn.datasets import cifar10
(X,Y),(X_test,Y_test)=cifar10.load_data("/dataset")
Y=to_categorical(Y,10)
Y_test=to_categorical(Y_test,10)
print(X.shape)
print(Y.shape)
print(X_test.shape)
print(Y_test.shape)

# # Create DASK array using numpy arrays
# import dask.array as da
# X = da.from_array(np.asarray(X), chunks=(1000))
# Y = da.from_array(np.asarray(Y), chunks=(1000))
#
# X_test = da.from_array(np.asarray(X_test), chunks=(1000))
# Y_test = da.from_array(np.asarray(Y_test), chunks=(1000))

# # Create a hdf5 dataset from CIFAR-10 numpy array
import h5py
h5f = h5py.File('data.h5', 'w')
h5f.create_dataset('cifar10_X', data=X)
h5f.create_dataset('cifar10_Y', data=Y)
h5f.create_dataset('cifar10_X_test', data=X_test)
h5f.create_dataset('cifar10_Y_test', data=Y_test)
h5f.close()

# Load hdf5 dataset
h5f = h5py.File('data.h5', 'r')
X = h5f['cifar10_X']
Y = h5f['cifar10_Y']
X_test = h5f['cifar10_X_test']
Y_test = h5f['cifar10_Y_test']

# Build network
network = input_data(shape=[None, 32, 32, 3])
network = conv_2d(network, 32, 3, activation='relu')
network = max_pool_2d(network, 2)
network = dropout(network, 0.75)
network = conv_2d(network, 64, 3, activation='relu')
network = conv_2d(network, 64, 3, activation='relu')
network = max_pool_2d(network, 2)
network = dropout(network, 0.5)
network = fully_connected(network, 512, activation='relu')
network = dropout(network, 0.5)
network = fully_connected(network, 10, activation='softmax')
network = regression(network, optimizer='adam',
                     loss='categorical_crossentropy',
                     learning_rate=0.001)

# Training
model = tflearn.DNN(network, tensorboard_verbose=0,max_checkpoints=30)
model.fit(X, Y, n_epoch=1, shuffle=True, validation_set=(X_test, Y_test),
          show_metric=True, batch_size=96, run_id='cifar10_cnn')
model.save('./cifar10_cnn')
h5f.close()


#coding:utf-8
import cv2
import numpy as np
import os
import sys
from tqdm import tqdm

datapath = "data_set/UIUC/PNGImages/TrainImages"
def path(cls , i):
    return "%s/%s%d.png" % (datapath , cls , i + 1)

SAMPLES = 400

def get_extract_detect():
    return cv2.xfeatures2d.SIFT_create() , cv2.xfeatures2d.SIFT_create()

# 获得FLANN匹配器的函数
def get_flann_matcher():
    flann_params = dict(algorithm = 1 , trees = 5)

    return cv2.FlannBasedMatcher(flann_params , {})

def get_bow_extractor(extract , flann):
    return cv2.BOWImgDescriptorExtractor(extract , flann)

def extract_sift(fn , extractor , detector):
    im = cv2.imread(fn , 0)
    if type(im) == type(None):
        print('read image failture!')
        sys.exit(1)
    # print(fn)
    # print('descriptors：', extractor.compute(im, detector.detect(im))[1])
    return extractor.compute(im , detector.detect(im))[1]

# 使用视觉词袋集计算一张图像的描述符
# 返回每个图像的bow码本，即直方图，大小为1*clusterNum
def bow_features(img , extractor_bow , detector):
    return extractor_bow.compute(img , detector.detect(img))

def car_detection():
    pos, neg = "pos-", "neg-"
    detect , extract = get_extract_detect()

    matcher = get_flann_matcher()
    # 有1000个词
    bow_kmeans_trainer = cv2.BOWKMeansTrainer(1000)
    extract_bow = cv2.BOWImgDescriptorExtractor(extract, matcher)

    print('adding features to trainner')
    # 将提取的特征描述符添加到词袋中
    for i in tqdm(range(SAMPLES)):
        try:
            bow_kmeans_trainer.add(extract_sift(path(pos, i) , extract , detect))
            bow_kmeans_trainer.add(extract_sift(path(neg, i) , extract , detect))
        except:
            continue

    # 进行聚类得到dictionary，也就是聚类的中心
    voc = bow_kmeans_trainer.cluster()
    # voc不能为空
    extract_bow.setVocabulary(voc)

    print('adding to train data')
    traindata, trainlabels = [], []
    for i in tqdm(range(SAMPLES)):
        try:
            traindata.extend(bow_features(cv2.imread(path(pos, i) , 0) , extract_bow , detect))
            trainlabels.append(1)
            traindata.extend(bow_features(cv2.imread(path(neg, i) , 0) , extract_bow , detect))
            trainlabels.append(-1)
        except:
            pass

    svm = cv2.ml.SVM_create()
    svm.setType(cv2.ml.SVM_C_SVC)
    svm.setGamma(1)
    # 此参数可以决定分类器的训练误差和预测误差，其值越大，误判的可能习惯越小，但训练精度会降低。另一方面，值太低可能会导致过拟合，从而使预测精度降低
    svm.setC(35)
    svm.setKernel(cv2.ml.SVM_RBF)
    svm.train(np.array(traindata), cv2.ml.ROW_SAMPLE, np.array(trainlabels))

    return svm , extract_bow

if __name__ == '__main__':
    print('hello detector!')


# -*- coding: utf-8 -*-
"""
Created on Fri Aug 25 18:25:32 2017

@author: Alex
"""

import requests, json
from bs4 import BeautifulSoup

class getPrimiaryJsonData():
    def __init__(self, songID):
        self.songID = songID

    #解析获得歌曲的基本信息
    def paeser_music_detail(self):
        #获得主页的html信息
        main_res = requests.get('http://music.163.com/song?id={}'.format(self.songID))
        main_res.encoding = 'utf-8'
        html_sample = main_res.text
        soup = BeautifulSoup(html_sample, 'html.parser')
        #获得歌曲名字
        music_name = soup.select('.tit')[0].text
        #获得歌曲的歌手
        music_art = soup.select('p')[0].text
        #获取歌曲的专辑信息
        music_club = soup.select('p')[1].text
        return music_name, music_art, music_club

    #机械获得歌曲的歌词信息
    def parser_music_lyric(self):
        #获得歌词的html信息
        lyric_html = requests.get('http://music.163.com/api/song/media?id={}'.format(self.songID))
        lyric_html.encoding = 'utf-8'
        #由requests实例转化为unicode实例
        html_lyric = lyric_html.text
        #将unicode实例转化为json实例
        lyric_jd = json.loads(html_lyric, encoding='utf-8')
        #获得歌词信息
        music_lyric = lyric_jd['lyric']
        return music_lyric


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import re
import itertools

from ..core import Backend
from ...expr.arithmetic import *
from ...expr.math import *
from ...expr.datetimes import *
from ...expr.strings import *
from ...expr.element import *
from ...expr.reduction import *
from ...expr.collections import *
from ...expr.merge import *
from ...utils import output
from ..errors import CompileError
from ..utils import refresh_dynamic
from ... import types

class Analyzer(Backend):
    """
    Analyzer is used before optimzing,
    which analyze some operation that is not supported for this execution backend.
    """

    def __init__(self, dag, traversed=None, on_sub=None):
        self._expr = dag.root
        self._dag = dag
        self._indexer = itertools.count(0)
        self._traversed = traversed or set()
        self._on_sub = on_sub

    def analyze(self):
        for node in self._iter():
            self._traversed.add(id(node))
            self._visit_node(node)

        return self._dag.root

    def _iter(self):
        for node in self._expr.traverse(top_down=True, unique=True,
                                        traversed=self._traversed):
            yield node

        while True:
            all_traversed = True
            for node in self._expr.traverse(top_down=True, unique=True):
                if id(node) not in self._traversed:
                    all_traversed = False
                    yield node
            if all_traversed:
                break

    def _visit_node(self, node):
        try:
            node.accept(self)
        except NotImplementedError:
            return

    def _sub(self, expr, sub, parents=None):
        self._dag.substitute(expr, sub, parents=parents)
        if self._on_sub:
            self._on_sub(expr, sub)

    def _parents(self, expr):
        return self._dag.successors(expr)

    def _visit_cut(self, expr):
        is_seq = isinstance(expr, SequenceExpr)
        kw = dict()
        if is_seq:
            kw['_data_type'] = expr.dtype
        else:
            kw['_value_type'] = expr.dtype

        conditions = []
        thens = []

        if expr.include_under:
            bin = expr.bins[0]
            if expr.right and not expr.include_lowest:
                conditions.append(expr.input <= bin)
            else:
                conditions.append(expr.input < bin)
            thens.append(expr.labels[0])
        for i, bin in enumerate(expr.bins[1:]):
            lower_bin = expr.bins[i]
            if not expr.right or (i == 0 and expr.include_lowest):
                condition = lower_bin <= expr.input
            else:
                condition = lower_bin < expr.input

            if expr.right:
                condition = (condition & (expr.input <= bin))
            else:
                condition = (condition & (expr.input < bin))

            conditions.append(condition)
            if expr.include_under:
                thens.append(expr.labels[i + 1])
            else:
                thens.append(expr.labels[i])
        if expr.include_over:
            bin = expr.bins[-1]
            if expr.right:
                conditions.append(bin < expr.input)
            else:
                conditions.append(bin <= expr.input)
            thens.append(expr.labels[-1])

        sub = Switch(_conditions=conditions, _thens=thens, **kw)
        self._sub(expr, sub)

    def visit_element_op(self, expr):
        if isinstance(expr, Between):
            if expr.inclusive:
                sub = ((expr.left <= expr.input) & (expr.input.copy() <= expr.right))
            else:
                sub = ((expr.left < expr.input) & (expr.input.copy() < expr.right))
            self._sub(expr, sub.rename(expr.name))
        elif isinstance(expr, Cut):
            self._visit_cut(expr)
        else:
            raise NotImplementedError

    def visit_sample(self, expr):
        if expr._parts is None:
            raise CompileError('ODPS SQL only support sampling by specifying `parts` arg')

        idxes = [None, ] if expr._i is None else expr._i

        condition = None
        for idx in idxes:
            inputs = [expr._parts]
            if idx is not None:
                new_val = idx.value + 1
                inputs.append(Scalar(_value=new_val, _value_type=idx.value_type))
            if expr._sampled_fields:
                inputs.extend(expr._sampled_fields)
            cond = MappedExpr(_inputs=inputs, _func='SAMPLE', _data_type=types.boolean)
            if condition is None:
                condition = cond
            else:
                condition |= cond
        sub = FilterCollectionExpr(_input=expr.input, _predicate=condition,
                                   _schema=expr.schema)
        expr.input.optimize_banned = True

        self._sub(expr, sub)

    def _visit_pivot(self, expr):
        columns_expr = expr.input.distinct([c.copy() for c in expr._columns])

        group_names = [g.name for g in expr._group]
        group_types = [g.dtype for g in expr._group]
        exprs = [expr]

        def callback(result, new_expr):
            expr = exprs[0]
            columns = [r[0] for r in result]

            if len(expr._values) > 1:
                names = group_names + \
                    ['{0}_{1}'.format(v.name, c)
                     for v in expr._values for c in columns]
                types = group_types + \
                        list(itertools.chain(*[[n.dtype] * len(columns)
                                               for n in expr._values]))
            else:
                names = group_names + columns
                types = group_types + [expr._values[0].dtype] * len(columns)
            new_expr._schema = Schema.from_lists(names, types)

            column_name = expr._columns[0].name  # column's size can only be 1
            values_names = [v.name for v in expr._values]

            @output(names, types)
            def reducer(keys):
                values = [None] * len(columns) * len(values_names)

                def h(row, done):
                    col = getattr(row, column_name)
                    for val_idx, value_name in enumerate(values_names):
                        val = getattr(row, value_name)
                        idx = len(columns) * val_idx + columns.index(col)
                        if values[idx] is not None:
                            raise ValueError('Row contains duplicate entries')
                        values[idx] = val
                    if done:
                        yield keys + tuple(values)

                return h

            fields = expr._group + expr._columns + expr._values
            pivoted = expr.input.select(fields).map_reduce(reducer=reducer, group=group_names)
            self._sub(new_expr, pivoted)

            # trigger refresh of dynamic operations
            refresh_dynamic(pivoted, self._dag)

        sub = CollectionExpr(_schema=DynamicSchema.from_lists(group_names, group_types),
                             _deps=[(columns_expr, callback)])
        self._sub(expr, sub)

    def _visit_pivot_table_without_columns(self, expr):
        def get_agg(field, agg_func, agg_func_name, fill_value):
            if isinstance(agg_func, six.string_types):
                aggregated = getattr(field, agg_func)()
            else:
                aggregated = field.agg(agg_func)
            if fill_value is not None:
                aggregated.fillna(fill_value)
            return aggregated.rename('{0}_{1}'.format(field.name, agg_func_name))

        grouped = expr.input.groupby(expr._group)
        aggs = []
        for agg_func, agg_func_name in zip(expr._agg_func, expr._agg_func_names):
            for value in expr._values:
                agg = get_agg(value, agg_func, agg_func_name, expr.fill_value)
                aggs.append(agg)
        sub = grouped.aggregate(aggs, sort_by_name=False)

        self._sub(expr, sub)

    def _visit_pivot_table_with_columns(self, expr):
        columns_expr = expr.input.distinct([c.copy() for c in expr._columns])

        group_names = [g.name for g in expr._group]
        group_types = [g.dtype for g in expr._group]
        exprs = [expr]

        def callback(result, new_expr):
            expr = exprs[0]
            columns = [r[0] for r in result]

            names = list(group_names)
            tps = list(group_types)
            aggs = []
            for agg_func_name, agg_func in zip(expr._agg_func_names, expr._agg_func):
                for value_col in expr._values:
                    for col in columns:
                        base = '{0}_'.format(col) if col is not None else ''
                        name = '{0}{1}_{2}'.format(base, value_col.name, agg_func_name)
                        names.append(name)
                        tps.append(value_col.dtype)

                        field = (expr._columns[0] == col).ifelse(
                            value_col, Scalar(_value_type=value_col.dtype))
                        if isinstance(agg_func, six.string_types):
                            agg = getattr(field, agg_func)()
                        else:
                            func = agg_func()

                            class ActualAgg(object):
                                def buffer(self):
                                    return func.buffer()

                                def __call__(self, buffer, value):
                                    if value is None:
                                        return
                                    func(buffer, value)

                                def merge(self, buffer, pbuffer):
                                    func.merge(buffer, pbuffer)

                                def getvalue(self, buffer):
                                    return func.getvalue(buffer)

                            agg = field.agg(ActualAgg)
                        if expr.fill_value is not None:
                            agg = agg.fillna(expr.fill_value)
                        agg = agg.rename(name)
                        aggs.append(agg)

            new_expr._schema = Schema.from_lists(names, tps)

            pivoted = expr.input.groupby(expr._group).aggregate(aggs, sort_by_name=False)
            self._sub(new_expr, pivoted)

            # trigger refresh of dynamic operations
            refresh_dynamic(pivoted, self._dag)

        sub = CollectionExpr(_schema=DynamicSchema.from_lists(group_names, group_types),
                             _deps=[(columns_expr, callback)])
        self._sub(expr, sub)

    def _visit_pivot_table(self, expr):
        if expr._columns is None:
            self._visit_pivot_table_without_columns(expr)
        else:
            self._visit_pivot_table_with_columns(expr)

    def visit_pivot(self, expr):
        if isinstance(expr, PivotCollectionExpr):
            self._visit_pivot(expr)
        else:
            self._visit_pivot_table(expr)

    def visit_value_counts(self, expr):
        collection = expr.input
        by = expr._by
        sort = expr._sort.value

        sub = collection.groupby(by).agg(count=by.count())
        if sort:
            sub = sub.sort('count', ascending=False)
        self._sub(expr, sub)

    def _gen_mapped_expr(self, expr, inputs, func, name,
                         args=None, kwargs=None, multiple=False):
        kwargs = dict(_inputs=inputs, _func=func, _name=name,
                      _func_args=args, _func_kwargs=kwargs,
                      _multiple=multiple)
        if isinstance(expr, SequenceExpr):
            kwargs['_data_type'] = expr.dtype
        else:
            kwargs['_value_type'] = expr.dtype
        return MappedExpr(**kwargs)

    def visit_binary_op(self, expr):
        if not options.df.analyze:
            raise NotImplementedError

        if isinstance(expr, FloorDivide):
            func = lambda l, r: l // r
            # multiple False will pass *args instead of namedtuple
            sub = self._gen_mapped_expr(expr, (expr.lhs, expr.rhs),
                                           func, expr.name, multiple=False)
            self._sub(expr, sub)
            return
        if isinstance(expr, Add) and \
                all(child.dtype == types.datetime for child in (expr.lhs, expr.rhs)):
            return
        elif isinstance(expr, (Add, Substract)):
            if expr.lhs.dtype == types.datetime and expr.rhs.dtype == types.datetime:
                pass
            elif any(isinstance(child, MilliSecondScalar) for child in (expr.lhs, expr.rhs)):
                pass
            else:
                return

            def func(l, r, method):
                from datetime import datetime, timedelta
                if not isinstance(l, datetime):
                    l = timedelta(milliseconds=l)
                if not isinstance(r, datetime):
                    r = timedelta(milliseconds=r)

                if method == '+':
                    res = l + r
                else:
                    res = l - r
                if isinstance(res, timedelta):
                    return int(res.microseconds / 1000)
                return res

            inputs = expr.lhs, expr.rhs, Scalar('+') if isinstance(expr, Add) else Scalar('-')
            sub = self._gen_mapped_expr(expr, inputs, func, expr.name, multiple=False)
            self._sub(expr, sub)

        raise NotImplementedError

    def visit_unary_op(self, expr):
        if not options.df.analyze:
            raise NotImplementedError

        if isinstance(expr, Invert) and isinstance(expr.input.dtype, types.Integer):
            sub = expr.input.map(lambda x: ~x)
            self._sub(expr, sub)
            return

        raise NotImplementedError

    def visit_math(self, expr):
        if not options.df.analyze:
            raise NotImplementedError

        if expr.dtype != types.decimal:
            if isinstance(expr, Arccosh):
                def func(x):
                    import numpy as np
                    return float(np.arccosh(x))
            elif isinstance(expr, Arcsinh):
                def func(x):
                    import numpy as np
                    return float(np.arcsinh(x))
            elif isinstance(expr, Arctanh):
                def func(x):
                    import numpy as np
                    return float(np.arctanh(x))
            elif isinstance(expr, Radians):
                def func(x):
                    import numpy as np
                    return float(np.radians(x))
            elif isinstance(expr, Degrees):
                def func(x):
                    import numpy as np
                    return float(np.degrees(x))
            else:
                raise NotImplementedError

            sub = expr.input.map(func, expr.dtype)
            self._sub(expr, sub)
            return

        raise NotImplementedError

    def visit_datetime_op(self, expr):
        if isinstance(expr, Strftime):
            if not options.df.analyze:
                raise NotImplementedError

            date_format = expr.date_format

            def func(x):
                return x.strftime(date_format)

            sub = expr.input.map(func, expr.dtype)
            self._sub(expr, sub)
            return

        raise NotImplementedError

    def visit_string_op(self, expr):
        if isinstance(expr, Ljust):
            rest = expr.width - expr.input.len()
            sub = expr.input + \
                     (rest >= 0).ifelse(expr._fillchar.repeat(rest), '')
            self._sub(expr, sub.rename(expr.name))
            return
        elif isinstance(expr, Rjust):
            rest = expr.width - expr.input.len()
            sub = (rest >= 0).ifelse(expr._fillchar.repeat(rest), '') + expr.input
            self._sub(expr, sub.rename(expr.name))
            return
        elif isinstance(expr, Zfill):
            fillchar = Scalar('0')
            rest = expr.width - expr.input.len()
            sub = (rest >= 0).ifelse(fillchar.repeat(rest), '') + expr.input
            self._sub(expr, sub.rename(expr.name))
            return

        if not options.df.analyze:
            raise NotImplementedError

        func = None
        if isinstance(expr, Contains) and (not expr.case or expr.flags > 0):
            flags = 0
            if not expr.case:
                flags = re.I
            if expr.flags > 0:
                flags = flags | expr.flags
            pat = expr.pat

            def func(x):
                r = re.compile(pat, flags)
                return r.search(x) is not None
        elif isinstance(expr, Extract) and expr.flags > 0:
            pat = expr.pat
            flags = expr.flags
            group = expr.group

            def func(x):
                r = re.compile(pat, flags)
                if group is None:
                    match = r.search(x)
                    if match:
                        return match.group()
                else:
                    match = r.search(x)
                    if match:
                        return match.group(group)
        elif isinstance(expr, Find) and expr.end is not None:
            start = expr.start
            end = expr.end
            substr = expr.sub

            def func(x):
                return x.find(substr, start, end)
        elif isinstance(expr, RFind):
            start = expr.start
            end = expr.end
            substr = expr.sub

            def func(x):
                return x.rfind(substr, start, end)
        elif isinstance(expr, (Lstrip, Strip, Rstrip)) and expr.to_strip != ' ':
            to_strip = expr.to_strip

            if isinstance(expr, Lstrip):
                def func(x):
                    return x.lstrip(to_strip)
            elif isinstance(expr, Strip):
                def func(x):
                    return x.strip(to_strip)
            elif isinstance(expr, Rstrip):
                def func(x):
                    return x.rstrip(to_strip)
        elif isinstance(expr, Pad):
            side = expr.side
            fillchar = expr.fillchar
            width = expr.width

            if side == 'left':
                func = lambda x: x.rjust(width, fillchar)
            elif side == 'right':
                func = lambda x: x.ljust(width, fillchar)
            elif side == 'both':
                func = lambda x: x.center(width, fillchar)
            else:
                raise NotImplementedError
        elif isinstance(expr, Slice):
            start, end, step = expr.start, expr.end, expr.step

            if end is None and step is None:
                raise NotImplementedError
            if isinstance(start, six.integer_types) and \
                    isinstance(end, six.integer_types) and step is None:
                if start >= 0 and end >= 0:
                    raise NotImplementedError

            has_start = start is not None
            has_end = end is not None
            has_step = step is not None

            def func(x, *args):
                idx = 0
                s, e, t = None, None, None
                for i in range(3):
                    if i == 0 and has_start:
                        s = args[idx]
                        idx += 1
                    if i == 1 and has_end:
                        e = args[idx]
                        idx += 1
                    if i == 2 and has_step:
                        t = args[idx]
                        idx += 1
                return x[s: e: t]

            inputs = expr.input, expr._start, expr._end, expr._step
            sub = self._gen_mapped_expr(expr, tuple(i for i in inputs if i is not None),
                                           func, expr.name, multiple=False)
            self._sub(expr, sub)
            return
        elif isinstance(expr, Swapcase):
            func = lambda x: x.swapcase()
        elif isinstance(expr, Title):
            func = lambda x: x.title()
        elif isinstance(expr, Strptime):
            date_format = expr.date_format

            def func(x):
                from datetime import datetime
                return datetime.strptime(x, date_format)
        else:
            if isinstance(expr, Isalnum):
                func = lambda x: x.isalnum()
            elif isinstance(expr, Isalpha):
                func = lambda x: x.isalpha()
            elif isinstance(expr, Isdigit):
                func = lambda x: x.isdigit()
            elif isinstance(expr, Isspace):
                func = lambda x: x.isspace()
            elif isinstance(expr, Islower):
                func = lambda x: x.islower()
            elif isinstance(expr, Isupper):
                func = lambda x: x.isupper()
            elif isinstance(expr, Istitle):
                func = lambda x: x.istitle()
            elif isinstance(expr, (Isnumeric, Isdecimal)):
                def u_safe(s):
                    try:
                        return unicode(s, "unicode_escape")
                    except:
                        return s

                if isinstance(expr, Isnumeric):
                    func = lambda x: u_safe(x).isnumeric()
                else:
                    func = lambda x: u_safe(x).isdecimal()

        if func is not None:
            sub = expr.input.map(func, expr.dtype)
            self._sub(expr, sub)
            return

        raise NotImplementedError

    def visit_reduction(self, expr):
        if isinstance(expr, (Var, GroupedVar)):
            sub = expr.input.std(ddof=expr._ddof) ** 2
            self._sub(expr, sub)

        raise NotImplementedError