"""
`Introduction <introyt1_tutorial.html>`_ ||
**Tensors** ||
`Autograd <autogradyt_tutorial.html>`_ ||
`Building Models <modelsyt_tutorial.html>`_ ||
`TensorBoard Support <tensorboardyt_tutorial.html>`_ ||
`Training Models <trainingyt_tutorial.html>`_ ||
`Model Understanding <captumyt_tutorial.html>`_

Introduction to PyTorch Tensors
===============================

Follow along with the video below or on `youtube <https://www.youtube.com/watch?v=r7QDUPb2dCM>`__.

.. raw:: html

Tensors are the central data abstraction in PyTorch. This interactive
notebook provides an in-depth introduction to the ``torch.Tensor``
class.

First things first, let’s import the PyTorch module. We’ll also add
Python’s math module to facilitate some of the examples.

"""

import torch
import math

#########################################################################
# Creating Tensors
# ----------------
#
# The simplest way to create a tensor is with the ``torch.empty()`` call:
#

x = torch.empty(3, 4)
print(type(x))
print(x)

##########################################################################
# Let’s unpack what we just did:
#
# -  We created a tensor using one of the numerous factory methods
#    attached to the ``torch`` module.
# -  The tensor itself is 2-dimensional, having 3 rows and 4 columns.
# -  The type of the object returned is ``torch.Tensor``, which is an
#    alias for ``torch.FloatTensor``; by default, PyTorch tensors are
#    populated with 32-bit floating point numbers. (More on data types
#    below.)
# -  You will probably see some random-looking values when printing your
#    tensor. The ``torch.empty()`` call allocates memory for the tensor,
#    but does not initialize it with any values - so what you’re seeing is
#    whatever was in memory at the time of allocation.
#
# A brief note about tensors and their number of dimensions, and
# terminology:
#
# -  You will sometimes see a 1-dimensional tensor called a
#    *vector.*
# -  Likewise, a 2-dimensional tensor is often referred to as a
#    *matrix.*
# -  Anything with more than two dimensions is generally just
#    called a tensor.
#
# More often than not, you’ll want to initialize your tensor with some
# value. Common cases are all zeros, all ones, or random values, and the
# ``torch`` module provides factory methods for all of these:
#

zeros = torch.zeros(2, 3)
print(zeros)

ones = torch.ones(2, 3)
print(ones)

torch.manual_seed(1729)
random = torch.rand(2, 3)
print(random)

#########################################################################
# The factory methods all do just what you’d expect - we have a tensor
# full of zeros, another full of ones, and another with random values
# between 0 and 1.
#
# Random Tensors and Seeding
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Speaking of the random tensor, did you notice the call to
# ``torch.manual_seed()`` immediately preceding it? Initializing tensors,
# such as a model’s learning weights, with random values is common but
# there are times - especially in research settings - where you’ll want
# some assurance of the reproducibility of your results. Manually setting
# your random number generator’s seed is the way to do this. Let’s look
# more closely:
#

torch.manual_seed(1729)
random1 = torch.rand(2, 3)
print(random1)

random2 = torch.rand(2, 3)
print(random2)

torch.manual_seed(1729)
random3 = torch.rand(2, 3)
print(random3)

random4 = torch.rand(2, 3)
print(random4)

############################################################################
# What you should see above is that ``random1`` and ``random3`` carry
# identical values, as do ``random2`` and ``random4``. Manually setting
# the RNG’s seed resets it, so that identical computations depending on
# random number should, in most settings, provide identical results.
#
# For more information, see the `PyTorch documentation on
# reproducibility <https://pytorch.org/docs/stable/notes/randomness.html>`__.
#
# Tensor Shapes
# ~~~~~~~~~~~~~
#
# Often, when you’re performing operations on two or more tensors, they
# will need to be of the same *shape* - that is, having the same number of
# dimensions and the same number of cells in each dimension. For that, we
# have the ``torch.*_like()`` methods:
#

x = torch.empty(2, 2, 3)
print(x.shape)
print(x)

empty_like_x = torch.empty_like(x)
print(empty_like_x.shape)
print(empty_like_x)

zeros_like_x = torch.zeros_like(x)
print(zeros_like_x.shape)
print(zeros_like_x)

ones_like_x = torch.ones_like(x)
print(ones_like_x.shape)
print(ones_like_x)

rand_like_x = torch.rand_like(x)
print(rand_like_x.shape)
print(rand_like_x)

#########################################################################
# The first new thing in the code cell above is the use of the ``.shape``
# property on a tensor. This property contains a list of the extent of
# each dimension of a tensor - in our case, ``x`` is a three-dimensional
# tensor with shape 2 x 2 x 3.
#
# Below that, we call the ``.empty_like()``, ``.zeros_like()``,
# ``.ones_like()``, and ``.rand_like()`` methods. Using the ``.shape``
# property, we can verify that each of these methods returns a tensor of
# identical dimensionality and extent.
#
# The last way to create a tensor that will cover is to specify its data
# directly from a PyTorch collection:
#

some_constants = torch.tensor([[3.1415926, 2.71828], [1.61803, 0.0072897]])
print(some_constants)

some_integers = torch.tensor((2, 3, 5, 7, 11, 13, 17, 19))
print(some_integers)

more_integers = torch.tensor(((2, 4, 6), [3, 6, 9]))
print(more_integers)

######################################################################
# Using ``torch.tensor()`` is the most straightforward way to create a
# tensor if you already have data in a Python tuple or list. As shown
# above, nesting the collections will result in a multi-dimensional
# tensor.
#
# .. note::
#      ``torch.tensor()`` creates a copy of the data.
#
# Tensor Data Types
# ~~~~~~~~~~~~~~~~~
#
# Setting the datatype of a tensor is possible a couple of ways:
#

a = torch.ones((2, 3), dtype=torch.int16)
print(a)

b = torch.rand((2, 3), dtype=torch.float64) * 20.
print(b)

c = b.to(torch.int32)
print(c)

##########################################################################
# The simplest way to set the underlying data type of a tensor is with an
# optional argument at creation time. In the first line of the cell above,
# we set ``dtype=torch.int16`` for the tensor ``a``. When we print ``a``,
# we can see that it’s full of ``1`` rather than ``1.`` - Python’s subtle
# cue that this is an integer type rather than floating point.
#
# Another thing to notice about printing ``a`` is that, unlike when we
# left ``dtype`` as the default (32-bit floating point), printing the
# tensor also specifies its ``dtype``.
#
# You may have also spotted that we went from specifying the tensor’s
# shape as a series of integer arguments, to grouping those arguments in a
# tuple. This is not strictly necessary - PyTorch will take a series of
# initial, unlabeled integer arguments as a tensor shape - but when adding
# the optional arguments, it can make your intent more readable.
#
# The other way to set the datatype is with the ``.to()`` method. In the
# cell above, we create a random floating point tensor ``b`` in the usual
# way. Following that, we create ``c`` by converting ``b`` to a 32-bit
# integer with the ``.to()`` method. Note that ``c`` contains all the same
# values as ``b``, but truncated to integers.
#
# Available data types include:
#
# -  ``torch.bool``
# -  ``torch.int8``
# -  ``torch.uint8``
# -  ``torch.int16``
# -  ``torch.int32``
# -  ``torch.int64``
# -  ``torch.half``
# -  ``torch.float``
# -  ``torch.double``
# -  ``torch.bfloat``
#
# Math & Logic with PyTorch Tensors
# ---------------------------------
#
# Now that you know some of the ways to create a tensor… what can you do
# with them?
#
# Let’s look at basic arithmetic first, and how tensors interact with
# simple scalars:
#

ones = torch.zeros(2, 2) + 1
twos = torch.ones(2, 2) * 2
threes = (torch.ones(2, 2) * 7 - 1) / 2
fours = twos ** 2
sqrt2s = twos ** 0.5

print(ones)
print(twos)
print(threes)
print(fours)
print(sqrt2s)

##########################################################################
# As you can see above, arithmetic operations between tensors and scalars,
# such as addition, subtraction, multiplication, division, and
# exponentiation are distributed over every element of the tensor. Because
# the output of such an operation will be a tensor, you can chain them
# together with the usual operator precedence rules, as in the line where
# we create ``threes``.
#
# Similar operations between two tensors also behave like you’d
# intuitively expect:
#

powers2 = twos ** torch.tensor([[1, 2], [3, 4]])
print(powers2)

fives = ones + fours
print(fives)

dozens = threes * fours
print(dozens)

##########################################################################
# It’s important to note here that all of the tensors in the previous code
# cell were of identical shape. What happens when we try to perform a
# binary operation on tensors if dissimilar shape?
#
# .. note::
#      The following cell throws a run-time error. This is intentional.
#
# ::
#
#    a = torch.rand(2, 3)
#    b = torch.rand(3, 2)
#
#    print(a * b)
#

##########################################################################
# In the general case, you cannot operate on tensors of different shape
# this way, even in a case like the cell above, where the tensors have an
# identical number of elements.
#
# In Brief: Tensor Broadcasting
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# .. note::
#      If you are familiar with broadcasting semantics in NumPy
#      ndarrays, you’ll find the same rules apply here.
#
# The exception to the same-shapes rule is *tensor broadcasting.* Here’s
# an example:
#

rand = torch.rand(2, 4)
doubled = rand * (torch.ones(1, 4) * 2)

print(rand)
print(doubled)

#########################################################################
# What’s the trick here? How is it we got to multiply a 2x4 tensor by a
# 1x4 tensor?
#
# Broadcasting is a way to perform an operation between tensors that have
# similarities in their shapes. In the example above, the one-row,
# four-column tensor is multiplied by *both rows* of the two-row,
# four-column tensor.
#
# This is an important operation in Deep Learning. The common example is
# multiplying a tensor of learning weights by a *batch* of input tensors,
# applying the operation to each instance in the batch separately, and
# returning a tensor of identical shape - just like our (2, 4) \* (1, 4)
# example above returned a tensor of shape (2, 4).
#
# The rules for broadcasting are:
#
# -  Each tensor must have at least one dimension - no empty tensors.
#
# -  Comparing the dimension sizes of the two tensors, *going from last to
#    first:*
#
#    -  Each dimension must be equal, *or*
#
#    -  One of the dimensions must be of size 1, *or*
#
#    -  The dimension does not exist in one of the tensors
#
# Tensors of identical shape, of course, are trivially “broadcastable”, as
# you saw earlier.
#
# Here are some examples of situations that honor the above rules and
# allow broadcasting:
#

a =     torch.ones(4, 3, 2)

b = a * torch.rand(   3, 2) # 3rd & 2nd dims identical to a, dim 1 absent
print(b)

c = a * torch.rand(   3, 1) # 3rd dim = 1, 2nd dim identical to a
print(c)

d = a * torch.rand(   1, 2) # 3rd dim identical to a, 2nd dim = 1
print(d)

#############################################################################
# Look closely at the values of each tensor above:
#
# -  The multiplication operation that created ``b`` was
#    broadcast over every “layer” of ``a``.
# -  For ``c``, the operation was broadcast over ever layer and row of
#    ``a`` - every 3-element column is identical.
# -  For ``d``, we switched it around - now every *row* is identical,
#    across layers and columns.
#
# For more information on broadcasting, see the `PyTorch
# documentation <https://pytorch.org/docs/stable/notes/broadcasting.html>`__
# on the topic.
#
# Here are some examples of attempts at broadcasting that will fail:
#
# .. note::
#      The following cell throws a run-time error. This is intentional.
#
# ::
#
#    a =     torch.ones(4, 3, 2)
#
#    b = a * torch.rand(4, 3)    # dimensions must match last-to-first
#
#    c = a * torch.rand(   2, 3) # both 3rd & 2nd dims different
#
#    d = a * torch.rand((0, ))   # can't broadcast with an empty tensor
#

###########################################################################
# More Math with Tensors
# ~~~~~~~~~~~~~~~~~~~~~~
#
# PyTorch tensors have over three hundred operations that can be performed
# on them.
#
# Here is a small sample from some of the major categories of operations:
#

# common functions
a = torch.rand(2, 4) * 2 - 1
print('Common functions:')
print(torch.abs(a))
print(torch.ceil(a))
print(torch.floor(a))
print(torch.clamp(a, -0.5, 0.5))

# trigonometric functions and their inverses
angles = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
sines = torch.sin(angles)
inverses = torch.asin(sines)
print('\nSine and arcsine:')
print(angles)
print(sines)
print(inverses)

# bitwise operations
print('\nBitwise XOR:')
b = torch.tensor([1, 5, 11])
c = torch.tensor([2, 7, 10])
print(torch.bitwise_xor(b, c))

# comparisons:
print('\nBroadcasted, element-wise equality comparison:')
d = torch.tensor([[1., 2.], [3., 4.]])
e = torch.ones(1, 2)  # many comparison ops support broadcasting!
print(torch.eq(d, e)) # returns a tensor of type bool

# reductions:
print('\nReduction ops:')
print(torch.max(d))        # returns a single-element tensor
print(torch.max(d).item()) # extracts the value from the returned tensor
print(torch.mean(d))       # average
print(torch.std(d))        # standard deviation
print(torch.prod(d))       # product of all numbers
print(torch.unique(torch.tensor([1, 2, 1, 2, 1, 2]))) # filter unique elements

# vector and linear algebra operations
v1 = torch.tensor([1., 0., 0.])         # x unit vector
v2 = torch.tensor([0., 1., 0.])         # y unit vector
m1 = torch.rand(2, 2)                   # random matrix
m2 = torch.tensor([[3., 0.], [0., 3.]]) # three times identity matrix

print('\nVectors & Matrices:')
print(torch.cross(v2, v1)) # negative of z unit vector (v1 x v2 == -v2 x v1)
print(m1)
m3 = torch.matmul(m1, m2)
print(m3)                  # 3 times m1
print(torch.svd(m3))       # singular value decomposition

##################################################################################
# This is a small sample of operations. For more details and the full inventory of
# math functions, have a look at the
# `documentation <https://pytorch.org/docs/stable/torch.html#math-operations>`__.
#
# Altering Tensors in Place
# ~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Most binary operations on tensors will return a third, new tensor. When
# we say ``c = a * b`` (where ``a`` and ``b`` are tensors), the new tensor
# ``c`` will occupy a region of memory distinct from the other tensors.
#
# There are times, though, that you may wish to alter a tensor in place -
# for example, if you’re doing an element-wise computation where you can
# discard intermediate values. For this, most of the math functions have a
# version with an appended underscore (``_``) that will alter a tensor in
# place.
#
# For example:
#

a = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('a:')
print(a)
print(torch.sin(a))   # this operation creates a new tensor in memory
print(a)              # a has not changed

b = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('\nb:')
print(b)
print(torch.sin_(b))  # note the underscore
print(b)              # b has changed

#######################################################################
# For arithmetic operations, there are functions that behave similarly:
#

a = torch.ones(2, 2)
b = torch.rand(2, 2)

print('Before:')
print(a)
print(b)
print('\nAfter adding:')
print(a.add_(b))
print(a)
print(b)
print('\nAfter multiplying')
print(b.mul_(b))
print(b)

##########################################################################
# Note that these in-place arithmetic functions are methods on the
# ``torch.Tensor`` object, not attached to the ``torch`` module like many
# other functions (e.g., ``torch.sin()``). As you can see from
# ``a.add_(b)``, *the calling tensor is the one that gets changed in
# place.*
#
# There is another option for placing the result of a computation in an
# existing, allocated tensor. Many of the methods and functions we’ve seen
# so far - including creation methods! - have an ``out`` argument that
# lets you specify a tensor to receive the output. If the ``out`` tensor
# is the correct shape and ``dtype``, this can happen without a new memory
# allocation:
#

a = torch.rand(2, 2)
b = torch.rand(2, 2)
c = torch.zeros(2, 2)
old_id = id(c)

print(c)
d = torch.matmul(a, b, out=c)
print(c)                # contents of c have changed

assert c is d           # test c & d are same object, not just containing equal values
assert id(c), old_id    # make sure that our new c is the same object as the old one

torch.rand(2, 2, out=c) # works for creation too!
print(c)                # c has changed again
assert id(c), old_id    # still the same object!

##########################################################################
# Copying Tensors
# ---------------
#
# As with any object in Python, assigning a tensor to a variable makes the
# variable a *label* of the tensor, and does not copy it. For example:
#

a = torch.ones(2, 2)
b = a

a[0][1] = 561  # we change a...
print(b)       # ...and b is also altered

######################################################################
# But what if you want a separate copy of the data to work on? The
# ``clone()`` method is there for you:
#

a = torch.ones(2, 2)
b = a.clone()

assert b is not a      # different objects in memory...
print(torch.eq(a, b))  # ...but still with the same contents!

a[0][1] = 561          # a changes...
print(b)               # ...but b is still all ones

#########################################################################
# **There is an important thing to be aware of when using ``clone()``.**
# If your source tensor has autograd, enabled then so will the clone.
# **This will be covered more deeply in the video on autograd,** but if
# you want the light version of the details, continue on.
#
# *In many cases, this will be what you want.* For example, if your model
# has multiple computation paths in its ``forward()`` method, and *both*
# the original tensor and its clone contribute to the model’s output, then
# to enable model learning you want autograd turned on for both tensors.
# If your source tensor has autograd enabled (which it generally will if
# it’s a set of learning weights or derived from a computation involving
# the weights), then you’ll get the result you want.
#
# On the other hand, if you’re doing a computation where *neither* the
# original tensor nor its clone need to track gradients, then as long as
# the source tensor has autograd turned off, you’re good to go.
#
# *There is a third case,* though: Imagine you’re performing a computation
# in your model’s ``forward()`` function, where gradients are turned on
# for everything by default, but you want to pull out some values
# mid-stream to generate some metrics. In this case, you *don’t* want the
# cloned copy of your source tensor to track gradients - performance is
# improved with autograd’s history tracking turned off. For this, you can
# use the ``.detach()`` method on the source tensor:
#

a = torch.rand(2, 2, requires_grad=True) # turn on autograd
print(a)

b = a.clone()
print(b)

c = a.detach().clone()
print(c)

print(a)

#########################################################################
# What’s happening here?
#
# -  We create ``a`` with ``requires_grad=True`` turned on. **We haven’t
#    covered this optional argument yet, but will during the unit on
#    autograd.**
# -  When we print ``a``, it informs us that the property
#    ``requires_grad=True`` - this means that autograd and computation
#    history tracking are turned on.
# -  We clone ``a`` and label it ``b``. When we print ``b``, we can see
#    that it’s tracking its computation history - it has inherited
#    ``a``\ ’s autograd settings, and added to the computation history.
# -  We clone ``a`` into ``c``, but we call ``detach()`` first.
# -  Printing ``c``, we see no computation history, and no
#    ``requires_grad=True``.
#
# The ``detach()`` method *detaches the tensor from its computation
# history.* It says, “do whatever comes next as if autograd was off.” It
# does this *without* changing ``a`` - you can see that when we print
# ``a`` again at the end, it retains its ``requires_grad=True`` property.
#
# Moving to GPU
# -------------
#
# One of the major advantages of PyTorch is its robust acceleration on
# CUDA-compatible Nvidia GPUs. (“CUDA” stands for *Compute Unified Device
# Architecture*, which is Nvidia’s platform for parallel computing.) So
# far, everything we’ve done has been on CPU. How do we move to the faster
# hardware?
#
# First, we should check whether a GPU is available, with the
# ``is_available()`` method.
#
# .. note::
#      If you do not have a CUDA-compatible GPU and CUDA drivers
#      installed, the executable cells in this section will not execute any
#      GPU-related code.
#

if torch.cuda.is_available():
    print('We have a GPU!')
else:
    print('Sorry, CPU only.')

##########################################################################
# Once we’ve determined that one or more GPUs is available, we need to put
# our data someplace where the GPU can see it. Your CPU does computation
# on data in your computer’s RAM. Your GPU has dedicated memory attached
# to it. Whenever you want to perform a computation on a device, you must
# move *all* the data needed for that computation to memory accessible by
# that device. (Colloquially, “moving the data to memory accessible by the
# GPU” is shorted to, “moving the data to the GPU”.)
#
# There are multiple ways to get your data onto your target device. You
# may do it at creation time:
#

if torch.cuda.is_available():
    gpu_rand = torch.rand(2, 2, device='cuda')
    print(gpu_rand)
else:
    print('Sorry, CPU only.')

##########################################################################
# By default, new tensors are created on the CPU, so we have to specify
# when we want to create our tensor on the GPU with the optional
# ``device`` argument. You can see when we print the new tensor, PyTorch
# informs us which device it’s on (if it’s not on CPU).
#
# You can query the number of GPUs with ``torch.cuda.device_count()``. If
# you have more than one GPU, you can specify them by index:
# ``device='cuda:0'``, ``device='cuda:1'``, etc.
#
# As a coding practice, specifying our devices everywhere with string
# constants is pretty fragile. In an ideal world, your code would perform
# robustly whether you’re on CPU or GPU hardware. You can do this by
# creating a device handle that can be passed to your tensors instead of a
# string:
#

if torch.cuda.is_available():
    my_device = torch.device('cuda')
else:
    my_device = torch.device('cpu')
print('Device: {}'.format(my_device))

x = torch.rand(2, 2, device=my_device)
print(x)

#########################################################################
# If you have an existing tensor living on one device, you can move it to
# another with the ``to()`` method. The following line of code creates a
# tensor on CPU, and moves it to whichever device handle you acquired in
# the previous cell.
#

y = torch.rand(2, 2)
y = y.to(my_device)

##########################################################################
# It is important to know that in order to do computation involving two or
# more tensors, *all of the tensors must be on the same device*. The
# following code will throw a runtime error, regardless of whether you
# have a GPU device available:
#
# ::
#
#    x = torch.rand(2, 2)
#    y = torch.rand(2, 2, device='gpu')
#    z = x + y  # exception will be thrown
#

###########################################################################
# Manipulating Tensor Shapes
# --------------------------
#
# Sometimes, you’ll need to change the shape of your tensor. Below, we’ll
# look at a few common cases, and how to handle them.
#
# Changing the Number of Dimensions
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# One case where you might need to change the number of dimensions is
# passing a single instance of input to your model. PyTorch models
# generally expect *batches* of input.
#
# For example, imagine having a model that works on 3 x 226 x 226 images -
# a 226-pixel square with 3 color channels. When you load and transform
# it, you’ll get a tensor of shape ``(3, 226, 226)``. Your model, though,
# is expecting input of shape ``(N, 3, 226, 226)``, where ``N`` is the
# number of images in the batch. So how do you make a batch of one?
#

a = torch.rand(3, 226, 226)
b = a.unsqueeze(0)

print(a.shape)
print(b.shape)

##########################################################################
# The ``unsqueeze()`` method adds a dimension of extent 1.
# ``unsqueeze(0)`` adds it as a new zeroth dimension - now you have a
# batch of one!
#
# So if that’s *un*\ squeezing? What do we mean by squeezing? We’re taking
# advantage of the fact that any dimension of extent 1 *does not* change
# the number of elements in the tensor.
#

c = torch.rand(1, 1, 1, 1, 1)
print(c)

##########################################################################
# Continuing the example above, let’s say the model’s output is a
# 20-element vector for each input. You would then expect the output to
# have shape ``(N, 20)``, where ``N`` is the number of instances in the
# input batch. That means that for our single-input batch, we’ll get an
# output of shape ``(1, 20)``.
#
# What if you want to do some *non-batched* computation with that output -
# something that’s just expecting a 20-element vector?
#

a = torch.rand(1, 20)
print(a.shape)
print(a)

b = a.squeeze(0)
print(b.shape)
print(b)

c = torch.rand(2, 2)
print(c.shape)

d = c.squeeze(0)
print(d.shape)

#########################################################################
# You can see from the shapes that our 2-dimensional tensor is now
# 1-dimensional, and if you look closely at the output of the cell above
# you’ll see that printing ``a`` shows an “extra” set of square brackets
# ``[]`` due to having an extra dimension.
#
# You may only ``squeeze()`` dimensions of extent 1. See above where we
# try to squeeze a dimension of size 2 in ``c``, and get back the same
# shape we started with. Calls to ``squeeze()`` and ``unsqueeze()`` can
# only act on dimensions of extent 1 because to do otherwise would change
# the number of elements in the tensor.
#
# Another place you might use ``unsqueeze()`` is to ease broadcasting.
# Recall the example above where we had the following code:
#
# ::
#
#    a =     torch.ones(4, 3, 2)
#
#    c = a * torch.rand(   3, 1) # 3rd dim = 1, 2nd dim identical to a
#    print(c)
#
# The net effect of that was to broadcast the operation over dimensions 0
# and 2, causing the random, 3 x 1 tensor to be multiplied element-wise by
# every 3-element column in ``a``.
#
# What if the random vector had just been 3-element vector? We’d lose the
# ability to do the broadcast, because the final dimensions would not
# match up according to the broadcasting rules. ``unsqueeze()`` comes to
# the rescue:
#

a = torch.ones(4, 3, 2)
b = torch.rand(   3)     # trying to multiply a * b will give a runtime error
c = b.unsqueeze(1)       # change to a 2-dimensional tensor, adding new dim at the end
print(c.shape)
print(a * c)             # broadcasting works again!

######################################################################
# The ``squeeze()`` and ``unsqueeze()`` methods also have in-place
# versions, ``squeeze_()`` and ``unsqueeze_()``:
#

batch_me = torch.rand(3, 226, 226)
print(batch_me.shape)
batch_me.unsqueeze_(0)
print(batch_me.shape)

##########################################################################
# Sometimes you’ll want to change the shape of a tensor more radically,
# while still preserving the number of elements and their contents. One
# case where this happens is at the interface between a convolutional
# layer of a model and a linear layer of the model - this is common in
# image classification models. A convolution kernel will yield an output
# tensor of shape *features x width x height,* but the following linear
# layer expects a 1-dimensional input. ``reshape()`` will do this for you,
# provided that the dimensions you request yield the same number of
# elements as the input tensor has:
#

output3d = torch.rand(6, 20, 20)
print(output3d.shape)

input1d = output3d.reshape(6 * 20 * 20)
print(input1d.shape)

# can also call it as a method on the torch module:
print(torch.reshape(output3d, (6 * 20 * 20,)).shape)

###############################################################################
# .. note::
#      The ``(6 * 20 * 20,)`` argument in the final line of the cell
#      above is because PyTorch expects a **tuple** when specifying a
#      tensor shape - but when the shape is the first argument of a method, it
#      lets us cheat and just use a series of integers. Here, we had to add the
#      parentheses and comma to convince the method that this is really a
#      one-element tuple.
#
# When it can, ``reshape()`` will return a *view* on the tensor to be
# changed - that is, a separate tensor object looking at the same
# underlying region of memory. *This is important:* That means any change
# made to the source tensor will be reflected in the view on that tensor,
# unless you ``clone()`` it.
#
# There *are* conditions, beyond the scope of this introduction, where
# ``reshape()`` has to return a tensor carrying a copy of the data. For
# more information, see the
# `docs <https://pytorch.org/docs/stable/torch.html#torch.reshape>`__.
#

#######################################################################
# NumPy Bridge
# ------------
#
# In the section above on broadcasting, it was mentioned that PyTorch’s
# broadcast semantics are compatible with NumPy’s - but the kinship
# between PyTorch and NumPy goes even deeper than that.
#
# If you have existing ML or scientific code with data stored in NumPy
# ndarrays, you may wish to express that same data as PyTorch tensors,
# whether to take advantage of PyTorch’s GPU acceleration, or its
# efficient abstractions for building ML models. It’s easy to switch
# between ndarrays and PyTorch tensors:
#

import numpy as np

numpy_array = np.ones((2, 3))
print(numpy_array)

pytorch_tensor = torch.from_numpy(numpy_array)
print(pytorch_tensor)

##########################################################################
# PyTorch creates a tensor of the same shape and containing the same data
# as the NumPy array, going so far as to keep NumPy’s default 64-bit float
# data type.
#
# The conversion can just as easily go the other way:
#

pytorch_rand = torch.rand(2, 3)
print(pytorch_rand)

numpy_rand = pytorch_rand.numpy()
print(numpy_rand)

##########################################################################
# It is important to know that these converted objects are using *the same
# underlying memory* as their source objects, meaning that changes to one
# are reflected in the other:
#

numpy_array[1, 1] = 23
print(pytorch_tensor)

pytorch_rand[1, 1] = 17
print(numpy_rand)


""" Pruning related functions.

Reference - https://jacobgil.github.io/deeplearning/pruning-deep-learning
"""
import os
import re
import time
import functools
import logging
import copy
from collections import namedtuple, OrderedDict, Counter
module_logger = logging.getLogger('pruning')
module_logger.setLevel(logging.INFO)

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.autograd import Variable

from taylor_pruning.transfer import ModelTransfer
from taylor_pruning.utils import *
from taylor_pruning.mask import ActMask

ActRank = namedtuple('ActRank', ['mod_name', 'channel', 'crit'])

def get_act_hook(mod, inp, out, mod_name=None, act_map=None):
  """ The forward hook to collect activation. """
  assert isinstance(mod_name, str)
  assert isinstance(act_map, dict)

  act_map[mod_name] = out

def get_grad_hook(mod, grad_in, grad_out, mod_name=None, grad_map=None):
  """ The hook to collect gradient. """
  assert isinstance(mod_name, str)
  assert isinstance(grad_map, dict)
  assert len(grad_out) == 1

  grad_map[mod_name] = grad_out[0]

def register_hooks(model, act_map, grad_map, logger=None):
  """ Register hooks on model to collect activations and
    gradients to act_map and grad_map respectively.

    Args:
      model(nn.Module):
      act_map(dict): activation map
      grad_map(dict): gradient map
      logger(Logger): log information
  """

  assert isinstance(model, nn.Module)
  assert isinstance(act_map, dict)
  assert isinstance(grad_map, dict)

  if logger is None:
    logger = module_logger

  hooks = {}

  # register hooks
  for name, mod in model.named_modules():
    # skip modules unnecessary for pruning
    if not isinstance(mod, nn.Conv2d) and not isinstance(mod, nn.Linear):
      continue

    hooks[name] = []
    hooks[name].append(
        mod.register_forward_hook(
            functools.partial(get_act_hook, act_map=act_map, mod_name=name)))
    logger.debug('Registered act hook for module "{}"'.format(name))
    hooks[name].append(
        mod.register_backward_hook(
            functools.partial(get_grad_hook, grad_map=grad_map, mod_name=name)))
    logger.debug('Registered grad hook for module "{}"'.format(name))

  return hooks

def compute_taylor_criterion(act, grad):
  """ Compute the taylor criterion Eq. (8) of a single module. """
  assert isinstance(act, torch.Tensor)
  assert isinstance(grad, torch.Tensor)
  assert act.shape == grad.shape

  # turn both into a tensor of (batch, channels, img_h, img_w)
  if len(act.shape) == 2:  # especially for nn.Linear
    act = act.view([*act.shape, 1, 1])
    grad = grad.view([*grad.shape, 1, 1])

  with torch.no_grad():  # this is critical for the memory usage
    crit = torch.mul(grad, act)  # element-wise
    crit = crit.view([*crit.shape[:2], -1])  # flatten spatial dims
    crit = crit.mean(dim=2)  # average spatially
    crit = torch.abs(crit)  # take abs for the final crit value
    crit = crit.mean(dim=0, keepdim=True)  # average across all mini-batches

  return crit

def rank_act(crit_map):
  """ Rank channels in activations by their criterion.

    NOTE: All the modules that are affected by replacement should be
      returned by this function.

    Args:
      crit_map(dict): computed criterion
    Returns:
      A list of ActRank(<mod_name, channel_id, crit_val>)
      sorted in increasing order
  """
  ranking = []

  for mod_name in crit_map:
    crit = crit_map[mod_name].detach().cpu().numpy()
    crit = np.squeeze(crit)
    assert len(crit.shape) == 1

    crit = crit.tolist()
    # generate the list of tuple for every module
    ranking.extend(list([ActRank(mod_name, i, c) for i, c in enumerate(crit)]))

  return list(sorted(ranking, key=lambda k: k.crit))  # sort by crit value

def create_crit_map(act_map, grad_map):
  """ Create crit_map from act_map and grad_map. """
  assert isinstance(act_map, dict)
  assert isinstance(grad_map, dict)

  # compute and store criterion
  crit_map = {}
  for key in grad_map:
    act, grad = act_map[key], grad_map[key]
    crit = compute_taylor_criterion(act, grad)
    crit = crit.detach().cpu()

    if key not in crit_map:
      crit_map[key] = crit
    else:
      crit_map[key] = torch.cat((crit_map[key], crit), dim=0)

  return crit_map

def get_mod_channels(mod, out=False):
  """ Collect the number of channels. """
  if isinstance(mod, nn.Conv2d):
    return mod.out_channels if out else mod.in_channels
  elif isinstance(mod, nn.Linear):
    return mod.out_features if out else mod.in_features
  else:
    raise TypeError('Cannot recognise module with type: {}'.format(type(mod)))

def update_mod_channels(mod, channels, out=False):
  if isinstance(mod, nn.Conv2d):
    if out:
      mod.out_channels = channels
    else:
      mod.in_channels = channels
  elif isinstance(mod, nn.Linear):
    if out:
      mod.out_features = channels
    else:
      mod.in_features = channels
  else:
    raise TypeError('Cannot recognise module with type: {}'.format(type(mod)))

def clone_module(mod, in_channels, out_channels):
  """ Create a new nn.Conv2d or nn.Linear module based on
    the original module provided. """
  if isinstance(mod, nn.Conv2d):
    mod_ = nn.Conv2d(
        in_channels,
        out_channels,
        mod.kernel_size,
        stride=mod.stride,
        padding=mod.padding,
        bias=mod.bias is not None,
        dilation=mod.dilation)
  elif isinstance(mod, nn.Linear):
    mod_ = nn.Linear(in_channels, out_channels, bias=mod.bias is not None)
  else:
    raise TypeError('Cannot recognise module with type: {}'.format(type(mod)))

  return mod_

def get_channels_to_prune(ranking,
                          num_channels,
                          mod_map=None,
                          par_map=None,
                          least_num_channels=8,
                          excludes=None):
  """ Return the channels that will be pruned.

  Args:
    ranking(list): a list of ActRank, already SORTED.
    num_channels(int): the amount of channels that will be pruned.
    mod_map(dict): mapping from module name to module
    par_map(dict): mapping from module name to its parent's name
    least_num_channels(int): at least leave each activation
      this amount of channels
    excludes(list): a list of mod names that will be excluded
      for pruning.
  Returns:
    A map from mod name to a list of channel indices
  """
  assert isinstance(least_num_channels, int) and least_num_channels >= 0

  # map from mod_name to channels
  rc_map = {}
  pc_map = {}

  # iterate rankings to initialise rc_map
  for rank in ranking:
    if rank.mod_name not in rc_map:
      rc_map[rank.mod_name] = set()
    rc_map[rank.mod_name].add(rank.channel)  # insert all channels appeared

    if rank.mod_name not in pc_map:
      pc_map[rank.mod_name] = set()

  # update rc_map by existing masks
  # when mod_map and par_map are provided
  if mod_map is not None and par_map is not None:
    for mod_name in rc_map:
      mod = mod_map[mod_name]
      par = mod_map[par_map[mod_name]]

      act_mask = find_act_mask(mod, par)
      # use act_mask to update the set of remaining channels
      if act_mask is not None:
        mask_val = act_mask.mask.cpu().numpy()
        rc_map[mod_name] = set(np.nonzero(mask_val)[0].tolist())
        pc_map[mod_name] = set(np.nonzero(mask_val == 0)[0].tolist())

  if excludes is None:
    excludes = []

  cnt = 0
  for rk in ranking:
    # this channel has been removed
    if rk.channel not in rc_map[rk.mod_name]:
      continue
    # remaining channels are less than the threshold
    if len(rc_map[rk.mod_name]) <= least_num_channels:
      continue
    # this module is excluded
    if rk.mod_name in excludes:
      continue

    # update
    rc_map[rk.mod_name].remove(rk.channel)
    pc_map[rk.mod_name].add(rk.channel)

    cnt += 1
    if cnt >= num_channels:
      break

  return pc_map

def find_act_mask(mod, par):
  """ Find the ActMask followed by mod in its parent children.

  Args:
    mod(nn.Module):
    par(nn.Module): parent module of mod
  Returns:
    The ActMask module or None
  """
  # locate mod in par.children
  mods = list(par.children())
  idx = find_in_parent(mod, par)  # locate the mod
  if idx == -1:
    raise ValueError('mod {} is not a children of {}.'.format(mod, par))
  if idx < len(mods) - 1 and isinstance(mods[idx + 1], ActMask):
    return mods[idx + 1]
  return None

def insert_or_update_act_mask(mod, par, ouc):
  """ Insert a new ActMask after mod or update its contents.

  Args:
    mod(nn.Module):
    par(nn.Module): parent module of mod
    ouc(list): a list of channel indices to be removed
  Returns:
    None
  """
  act_mask = find_act_mask(mod, par)

  # create the mask module
  if act_mask is None:
    num_channels = get_mod_channels(mod, out=True)
    act_mask = ActMask(num_channels)
    insert_after(act_mask, mod, par, name_suffix='mask')

  # update the mask value
  act_mask.mask.data[ouc] = 0.0

def prune_by_taylor_criterion(model,
                              crit_map,
                              num_channels_per_prune=1,
                              logger=None):
  """ Prune the given model by the crit_map.

  NOTE: need update for batch norm
  NOTE: need update for non-sequential connectivity

  Args:
    model(nn.Module): model to be pruned
    crit_map(dict): the computed criterion map
    num_channels_per_prune(int): number of channels will be pruned per
      iteration
  Returns:
    the model, updated in-place
  """
  assert isinstance(model, nn.Module)
  assert isinstance(crit_map, dict)
  assert isinstance(num_channels_per_prune, int)
  assert num_channels_per_prune >= 1

  if logger is None:
    logger = module_logger

  # build mod_map and par_map
  mod_map = OrderedDict()
  par_map = OrderedDict()  # record the parent module name
  for name, mod in model.named_modules():
    mod_map[name] = mod
    for child_name, child in mod.named_children():
      full_name = child_name if not name else '{}.{}'.format(name, child_name)
      par_map[full_name] = name

  # compute the ranking
  ranking = rank_act(crit_map)
  cls_name = list(model.named_modules())[-1][0]  # name of the classifier module
  pc_map = get_channels_to_prune(
      ranking,
      num_channels_per_prune,
      mod_map=mod_map,
      par_map=par_map,
      excludes=[cls_name])

  # iterate every activation
  # in each iteration, we remove those channels that
  # are marked to be pruned, by creating a new Conv2d
  # or Linear model and migrate the weights to the new
  # model.
  for i, (name, channels_to_prune) in enumerate(pc_map.items()):
    mod = mod_map[name]
    assert isinstance(mod, nn.Conv2d) or isinstance(mod, nn.Linear)

    insert_or_update_act_mask(mod, mod_map[par_map[name]],
                              list(channels_to_prune))

    # NOTE: this is an older module replacement approach. We now prefer
    #   adding mask modules
    # collect the number of channels from the module
    # that outputs the current activation.
    # in_channels = get_mod_channels(mod, out=False)
    # out_channels = get_mod_channels(mod, out=True)
    #
    # pre = mod_map[act_chl[i - 1][0]] if i > 0 else None
    # ouc = [x for x in range(out_channels) if x not in act_chl[i][1]]
    # if i == 0:
    #   inc = range(in_channels)
    # else:
    #   # HACK: the interface between CONV-FC is hard to handle
    #   if isinstance(mod, nn.Linear) and isinstance(pre, nn.Conv2d):
    #     img_size = in_channels // mod_out[act_chl[i - 1][0]]
    #     inc = np.arange(in_channels).reshape((-1, img_size))
    #     inc_ = []
    #     for x in range(inc.shape[0]):
    #       if x not in act_chl[i - 1][1]:
    #         inc_.extend(inc[x, :].tolist())
    #     inc = inc_
    #   else:
    #     inc = [x for x in range(in_channels) if x not in act_chl[i - 1][1]]
    # mod_ = clone_module(mod, len(inc), len(ouc))
    # mod_.weight.data = mod.weight[ouc, :][:, inc]  # pruned weights
    # if mod_.bias is not None:
    #   mod_.bias.data = mod.bias[ouc]  # pruned biases
    # par_mod = mod_map[par_map[name]]
    # par_mod._modules[name.split('.')[-1]] = mod_
    # del mod  # explicitly remove this replaced module

  return model

class ModelPruner(ModelTransfer):
  """ Implemented pruning related utilities. """

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

    self.base_dir = args.checkpoint
    if isinstance(self.base_dir, str):
      os.makedirs(self.base_dir, exist_ok=True)

  def update_criterion(self, crit_map, act_map, grad_map, hooks):
    """ Called after forward-backward in one pass.

      act_map and grad_map will be recycled
      hooks in `hooks` will be removed
    """
    assert isinstance(crit_map, dict)
    assert isinstance(act_map, dict)
    assert isinstance(grad_map, dict)
    assert isinstance(hooks, dict)  # NOTE: returned from register_hooks

    # compute and store criterion
    crit_map.update(create_crit_map(act_map, grad_map))

    # remove hooks
    for hooks_ in hooks.values():
      for hook in hooks_:
        hook.remove()

    del grad_map
    del act_map

  def compute_taylor_criterion(self, model, use_cuda=True, logger=None):
    """ Compute the taylor criterion for every activation in the model. """
    assert isinstance(model, nn.Module)

    if logger is None:
      logger = module_logger

    batch_time = AverageMeter()
    data_time = AverageMeter()
    crit_time = AverageMeter()

    # switch to train mode
    model.train()

    crit_map = {}

    end = time.time()
    for i, (x, target) in enumerate(self.train_loader):
      data_time.update(time.time() - end)

      # create hooks
      grad_map, act_map = {}, {}
      hooks = register_hooks(model, act_map, grad_map)

      if use_cuda:
        x = x.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)

      y = model(x)
      loss = self.criterion(y, target)
      loss.backward()

      # collect criterion
      # NOTE: here is the most time-consuming part (or maybe not)
      crit_start = time.time()
      self.update_criterion(crit_map, act_map, grad_map, hooks)
      crit_time.update(time.time() - crit_start)

      batch_time.update(time.time() - end)
      end = time.time()

      if i % self.args.print_freq == 0:
        print('[{0:3d}/{1:3d}]\t'
              'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
              'Crit {crit_time.val:.3f} ({crit_time.avg:.3f})\t'
              'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'.format(
                  i,
                  len(self.train_loader),
                  crit_time=crit_time,
                  batch_time=batch_time,
                  data_time=data_time))

    # summarise result
    logger.info('==> Summarising criterion from all batches ...')
    for key in crit_map:
      crit_map[key] = torch.mean(crit_map[key], dim=0)

    return crit_map

  def prune_loop(self, model, use_cuda=True, logger=None):
    """ Prune for the required number of iterations. """
    if logger is None:
      logger = module_logger

    total_num_channels_pruned = 0

    for prune_iter in range(self.args.num_prune_iters):
      logger.info('==> Running pruning iteration [{:3d}/{:3d}] ...'.format(
          prune_iter, self.args.num_prune_iters))

      # create the checkpoint directory
      checkpoint = os.path.join(self.base_dir,
                                'prune_iter_{}'.format(prune_iter))
      os.makedirs(checkpoint, exist_ok=True)

      # declare the new checkpoint
      self.args.checkpoint = checkpoint
      # NOTE: this logger is only used for the internal logic of training
      self.logger = self.get_logger(self.args)

      model = self.prune(
          model,
          total_num_channels_pruned=total_num_channels_pruned,
          num_channels_per_prune=self.args.num_channels_per_prune,
          use_cuda=use_cuda,
          logger=logger)
      total_num_channels_pruned += self.args.num_channels_per_prune

  def prune(self,
            model,
            total_num_channels_pruned=0,
            num_channels_per_prune=0,
            use_cuda=True,
            logger=None):
    """ A prune round. """
    if logger is None:
      logger = module_logger

    # collect the criterion map
    logger.info('==> Collecting criterion for all modules ...')
    crit_map = self.compute_taylor_criterion(
        model, use_cuda=use_cuda, logger=logger)

    # prune the model based on the criterion
    logger.info('==> Pruning the model ...')
    model = prune_by_taylor_criterion(
        model,
        crit_map,
        num_channels_per_prune=num_channels_per_prune,
        logger=logger)

    total_num_channels_pruned += num_channels_per_prune
    logger.info('==> Pruned model')
    print(model)
    print(self.get_mask_status(model))
    self.sanity_check(model, total_num_channels_pruned)

    # train and validate this model
    logger.info('==> Fine-tuning the model ...')
    self.train(model, load_optim=False)
    self.validate(model)

    return model

  def sanity_check(self, model, total_num_channels_pruned):
    """ Check whether pruning runs correctly. """
    df = self.get_mask_status(model)

    if df['num_pruned'].sum() != total_num_channels_pruned:
      raise RuntimeError(
          'Number of pruned channels {} does not equal to required {}'.format(
              df['num_pruned'].sum(), total_num_channels_pruned))

  def get_mask_status(self, model):
    """ Print the status of each mask. """
    cols = ['name', 'num_pruned']
    data = []
    for name, mod in model.named_modules():
      if isinstance(mod, ActMask):
        data.append([name, np.count_nonzero(mod.mask.cpu().numpy() == 0)])

    return pd.DataFrame(data, columns=cols)


import csv
import os

import numpy as np
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler

# Here in this file, you should define functions to try out different encodings.
# Some options:
#   1) Bag of words. You don't need any library for this, it's simple enough to
#       implement on your own.
#   2) Word embeddings. You can use spacy or Word2Vec (or others, but these are good
#       starting points). Spacy will give you better embeddings, since they are
#       already defined. Word2Vec will create embeddings based on your dataset.

## Document the choices you make, including if you pre-process words/tokens by
# stemming, using POS (parts of speech info) or anything else.

## Create your own files to define Logistic regression/ Neural networks to try
# out the performace of different ML algorithms. It'll be up to you to evaluate
# the performance.

class SentimentDataset(Dataset):
    """SentimentDataset [summary]

    [extended_summary]

    :param path_to_data: Path to dataset directory
    :type path_to_data: str
    """
    def __init__(self, path_to_data):
        ## TODO: Initialise the dataset given the path to the dataset directory.
        ## You may want to include other parameters, totally your choice.
        pass

    def __len__(self):
        """__len__ [summary]

        [extended_summary]
        """
        ## TODO: Returns the length of the dataset.
        pass

    def __getitem__(self, index):
        """__getitem__ [summary]

        [extended_summary]

        :param index: [description]
        :type index: [type]
        """
        ## TODO: This returns only ONE sample from the dataset, for a given index.
        ## The returned sample should be a tuple (x, y) where x is your input
        ## vector and y is your label
        ## Before returning your sample, you should check if there is a transform
        ## sepcified, and pply that transform to your sample
        # Eg:
        # if self.transform:
        #   sample = self.transform(sample)
        ## Remember to convert the x and y into torch tensors.

        pass

def get_data_loaders(path_to_pkl,
                     path_to_labels,
                     train_val_test=[0.8, 0.2, 0.2],
                     batch_size=32):
  """
  You know the drill by now.
  """
  pass


# normal libraries
import unittest
from unittest import TestCase

import numpy as np
import matplotlib.pyplot as plt

from PIL import ImageChops, Image  # image comparison

import os
import pathlib
# other files
from corai_plot import APlot
from config import ROOT_DIR
from corai_metaclass import Register, deco_register, dict_register_classes

"""
    Auto
    False - manual testing
          - save the file and check manually if the behaviour is as expected
    True - automatic testing
         - save the file with test_name and compare against manually approved correct image
    ! Only use false locally
"""

AUTO = True
DELETE_TEST_PLOT = True
PATH = os.path.join(ROOT_DIR, 'corai_plot', 'tests', 'image_reference_test_plot')

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def image_comparison(path1, path2):
    im1 = Image.open(path1).convert('RGB')  # we open the image with RGB code,
    # because some times there is noise on the image
    im2 = Image.open(path2).convert('RGB')  # same
    diff = ImageChops.difference(im2, im1)
    if diff.getbbox():  # if images are different
        return False
    else:
        return True

class Test_APlot(TestCase):
    def setUp(self) -> None:
        self.xx = np.linspace(0, 10, 10000)
        self.yy = np.cos(self.xx)
        self.image_name = ""
        return

    def tearDown(self):
        if DELETE_TEST_PLOT:
            for file in os.listdir(PATH):
                if file.startswith("test") and file.endswith(".png"):
                    file_path = os.path.join(PATH, file)
                    os.remove(file_path)

    def check_plot(self):
        image_path = os.path.join(PATH, f"{self.image_name}.png")
        if AUTO:
            test_path = os.path.join(PATH, f"test_{self.image_name}.png")
            plt.savefig(test_path, dpi=50)
            assert (image_comparison(test_path, image_path))
        else:
            plt.savefig(image_path, dpi=50)

    def test_constructor_plot_data_directly_only_yy(self):
        # direct plot only yy
        APlot(datay=self.yy)
        self.image_name = "image_yy"
        self.check_plot()

    def test_constructor_plot_data_directly_xx_yy(self):
        # direct plot xx and yy
        APlot(datax=self.xx, datay=self.yy)
        self.image_name = "image_xx_yy"
        self.check_plot()

    def test_set_dict_ax_and_bis_each_parameter_only_first_two_axis_same_graph(self):
        # first, trying every simple possibility for dict_ax
        yy = np.abs(np.cos(self.xx)) + 1
        with self.subTest('title'):
            aplot = APlot()
            dict_plot = {'title': 'my title'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_title"
            self.check_plot()

        with self.subTest('xlabel'):
            aplot = APlot()

            dict_plot = {'xlabel': 'my x label'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_xlabel"
            self.check_plot()

        with self.subTest('ylabel'):
            aplot = APlot()

            dict_plot = {'ylabel': 'my y label'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_ylabel"
            self.check_plot()

        with self.subTest('xscale'):
            aplot = APlot()

            dict_plot = {'xscale': 'log'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)
            self.image_name = "image_set_dict_ax_xscale_only_principal_axis"
            self.check_plot()

        with self.subTest('yscale left'):
            aplot = APlot()

            dict_plot = {'yscale': 'log'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_yscale"
            self.check_plot()

        with self.subTest('xint'):
            aplot = APlot()

            dict_plot = {'xint': True}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 1/3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_xint"
            self.check_plot()

        with self.subTest('xint_yint'):
            aplot = APlot()

            dict_plot = {'xint': True, 'yint': True}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 1/3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_xyint"
            self.check_plot()

        for i in range(11):
            with self.subTest('parameter', i=i):
                aplot = APlot()

                dict_plot = {'parameters': ['A', 3, 5]*i, 'name_parameters': ['A', '$\sigma$', '$\\rho$']*i}

                aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
                aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

                self.image_name = f"image_set_dict_ax_parameter_{i}"
                self.check_plot()

        # with self.subTest('parameters not good length'):
        #     aplot = APlot()
        #
        #     dict_plot = {
        #         'parameters': ['A', 3, 5, 10, 42], 'name_parameters': ['A']}
        #
        #     aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
        #     aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

        with self.subTest('xlim'):
            aplot = APlot()

            dict_plot = {'xlim': [0, 0.5]}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_xlim"
            self.check_plot()

        with self.subTest('ylim'):
            aplot = APlot()

            dict_plot = {'ylim': [1.2, 2.4]}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1)

            self.image_name = "image_set_dict_ax_ylim"
            self.check_plot()

    def test_set_dict_ax_and_bis_each_parameter_both_two_axis_same_graph(self):
        yy = np.abs(np.cos(self.xx)) + 1
        # second, try the same as before where the second axis is also setting another characteristic.
        with self.subTest('two different titles'):
            aplot = APlot()
            dict_plot1 = {'title': 'my title1'}
            dict_plot2 = {'title': 'my title2'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot1)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_title_2"
            self.check_plot()

        with self.subTest('two xlabels'):
            aplot = APlot()

            dict_plot1 = {'xlabel': 'my x label1'}
            dict_plot2 = {'xlabel': 'my x label2'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot1)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_xlabel_2"
            self.check_plot()

        with self.subTest('two ylabels'):
            aplot = APlot()

            dict_plot1 = {'ylabel': 'my y label1'}
            dict_plot2 = {'ylabel': 'my y label2'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot1)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_ylabel_2"
            self.check_plot()

        with self.subTest('xscale same both axis'):
            aplot = APlot()

            dict_plot = {'xscale': 'log'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot)

            self.image_name = "image_set_dict_ax_xscale_both_same_scale"
            self.check_plot()

        with self.subTest('yscale same both axis'):
            aplot = APlot()

            dict_plot = {'yscale': 'log'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot)

            self.image_name = "image_set_dict_ax_yscale_2_same"
            self.check_plot()

        with self.subTest('xscale different'):
            aplot = APlot()

            dict_plot2 = {'xscale': 'log'}
            dict_plot = {'xscale': 'linear'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_xscale_2_different_scales"
            self.check_plot()

        with self.subTest('yscale different'):
            aplot = APlot()

            dict_plot = {'yscale': 'log'}
            dict_plot2 = {'yscale': 'linear'}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_yscale_2_different"
            self.check_plot()

        # with self.subTest('xint same'):
        #     aplot = APlot()
        #
        #     dict_plot = {'xint': True}
        #
        #     aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
        #     aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot)
        #
        #     plt.savefig("image_reference_test_plot/test_image_set_dict_ax_xint_2_same.png")
        #
        # with self.subTest('xint_yint same'):
        #     aplot = APlot()
        #
        #     dict_plot = {'xint': True, 'yint': True}
        #
        #     aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
        #     aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot)
        #
        #     plt.savefig("image_reference_test_plot/test_image_set_dict_ax_xyint_same_2.png")

        # with self.subTest('xint different'):
        #     aplot = APlot()
        #
        #     dict_plot = {'xint': True}
        #     dict_plot2 = {'xint': False}
        #
        #     aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
        #     aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)
        #
        #     plt.savefig("image_reference_test_plot/test_image_set_dict_ax_xint_2_different_2.png")
        #
        # with self.subTest('xint_yint different'):
        #     aplot = APlot()
        #
        #     dict_plot = {'xint': True, 'yint': True}
        #     dict_plot2 = {'xint': True, 'yint': False}
        #
        #     aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot)
        #     aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)
        #
        #     plt.savefig("image_reference_test_plot/test_image_set_dict_ax_xyint_different_2.png")

        with self.subTest('xlim'):
            aplot = APlot()

            dict_plot1 = {'xlim': [0, 0.5]}
            dict_plot2 = {'xlim': [0, 1.5]}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot1)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_xlim_2"
            self.check_plot()

        with self.subTest('ylim'):
            aplot = APlot()

            dict_plot1 = {'ylim': [1, 5]}
            dict_plot2 = {'ylim': [1, 15]}

            aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot1)
            aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)

            self.image_name = "image_set_dict_ax_ylim_2"
            self.check_plot()

    def test_show_legend_for_two_axis(self):
        yy = np.abs(np.cos(self.xx)) + 1
        # second, try the same as before where the second axis is also setting another characteristic.

        aplot = APlot()
        dict_plot1 = {'title': 'my title1'}
        dict_plot2 = {'title': 'my title2'}

        aplot.uni_plot(0, self.xx, yy, dict_ax=dict_plot1)
        aplot.uni_plot_ax_bis(0, self.xx, 3 * np.sin(self.xx) + 1, dict_ax=dict_plot2)
        aplot.show_legend()
        self.image_name = "image_legend_two_axis"
        self.check_plot()

    def test_tight_layout(self):
        pass

    def test_uni_plot_one_plot(self):
        # uni plot
        aplot_2 = APlot()
        aplot_2.uni_plot(0, self.xx, self.yy)
        aplot_2.tight_layout()

        self.image_name = "image_uniplot_1"
        self.check_plot()

    def test_uni_plot_for_plots_same_graph(self):
        # uni plot *4
        aplot_3 = APlot(how=(2, 2))
        aplot_3.uni_plot(0, self.xx, self.yy)
        aplot_3.uni_plot(1, self.xx, self.yy)
        aplot_3.uni_plot(2, self.xx, self.yy)
        aplot_3.uni_plot(3, self.xx, self.yy)
        aplot_3.tight_layout()

        self.image_name = "image_uniplot_2"
        self.check_plot()

    def test_uni_plot_for_plots_same_graph_sharex(self):
        # uni plot *4
        aplot_3 = APlot(how=(2, 2), sharex=True)
        aplot_3.uni_plot(0, self.xx, self.yy)
        aplot_3.uni_plot(1, self.xx, self.yy)
        aplot_3.uni_plot(2, self.xx, self.yy)
        aplot_3.uni_plot(3, self.xx, self.yy)
        aplot_3.tight_layout()

        self.image_name = "image_uniplot_2_sharex"
        self.check_plot()

    def test_uni_plot_for_plots_same_graph_sharex_not_same_interval_for_x(self):
        # uni plot *4
        xx = np.linspace(-1, 1, 1000)
        aplot_3 = APlot(how=(2, 2), sharex=True)
        aplot_3.uni_plot(0, xx, np.exp(xx))
        aplot_3.uni_plot(1, xx + 1, np.exp(xx + 1))
        aplot_3.uni_plot(2, xx + 1, np.exp(xx + 1))
        aplot_3.uni_plot(3, xx, np.exp(xx))
        aplot_3.tight_layout()

        self.image_name = "image_uniplot_2_sharex_not_same_interval_for_x"
        self.check_plot()

    def test_uni_plot_for_plots_same_graph_sharey(self):
        # uni plot *4
        aplot_3 = APlot(how=(2, 2), sharey=True)
        aplot_3.uni_plot(0, self.xx, self.yy)
        aplot_3.uni_plot(1, self.xx, self.yy)
        aplot_3.uni_plot(2, self.xx, self.yy)
        aplot_3.uni_plot(3, self.xx, self.yy)
        aplot_3.tight_layout()

        self.image_name = "image_uniplot_2_sharey"
        self.check_plot()

    def test_uni_plot_for_plots_same_graph_sharey_not_same_interval_for_y(self):
        xx = np.linspace(-1, 1, 1000)
        aplot_3 = APlot(how=(2, 2), sharey=True)
        aplot_3.uni_plot(0, xx, np.exp(xx))
        aplot_3.uni_plot(1, xx + 1, np.exp(xx + 1))
        aplot_3.uni_plot(2, xx + 1, np.exp(xx + 1))
        aplot_3.uni_plot(3, xx, np.exp(xx))
        aplot_3.tight_layout()

        self.image_name = "image_uniplot_2_sharey_not_same_interval_for_y"
        self.check_plot()

    def test_bi_plot(self):
        aplot_4 = APlot(how=(2, 2))
        aplot_4.bi_plot(0, 1, self.xx, self.yy, self.xx, self.yy)
        aplot_4.uni_plot(2, self.xx, self.yy)
        aplot_4.uni_plot(3, self.xx, self.yy)
        aplot_4.tight_layout()

        self.image_name = "image_biplot"
        self.check_plot()

    def test_uni_plot_ax_bis(self):
        # two plots same figures
        aplot_1 = APlot()
        aplot_1.uni_plot(0, self.xx, self.yy + 5)
        aplot_1.uni_plot_ax_bis(0, self.xx, np.exp(self.xx))

        self.image_name = "image_plot_bis"
        self.check_plot()

    def test_cumulative_plot(self):
        aplot = APlot()
        values, base, _ = plt.hist(self.yy, bins=len(self.yy))
        aplot.cumulative_plot(base, values, nb_ax=0)

        self.image_name = "image_cumulative_plot"
        self.check_plot()
        pass

    def test_hist(self):
        aplot = APlot()
        aplot.hist(self.yy)
        self.image_name = "plot_hist"
        self.check_plot()

    def test_plot_function(self):
        def f(x):
            return 3 * x + 2 * x * x

        aplot = APlot()
        aplot.plot_function(f, self.xx)
        self.image_name = "image_plot_function"
        self.check_plot()

    def test_plot_vertical_line(self):
        aplot = APlot()
        aplot.plot_vertical_line(x=1, yy=np.array([2, 4]))

        self.image_name = "image_vertical_line"
        self.check_plot()

    def test_plot_line(self):
        aplot = APlot()
        aplot.plot_line(a=1, b=2, xx=self.xx)

        self.image_name = "image_line"
        self.check_plot()

    def test_plot_point(self):
        aplot = APlot()
        aplot.plot_point(x=1, y=1, dict_plot_param={'markersize': 5})

        self.image_name = "image_point"
        self.check_plot()

    def test_help_dict_plot(self):
        pass

    def test_help_dict_ax(self):
        pass


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.9.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + nbsphinx="hidden"
# %matplotlib inline

# + nbsphinx="hidden"
# %run notebook_setup
# -

# # Walkthrough: Examining DSHARP AS 209 Weights and Exporting Visibilities
#
# In this walkthrough tutorial, we'll use CASA tools to examine the visibilities, visibility residuals, and weights of a real multi-configuration dataset from the DSHARP survey.
#
# ## Obtaining and CLEANing the AS 209 measurement set
#
# The calibrated measurement sets from the DSHARP data release are available [online](https://almascience.eso.org/almadata/lp/DSHARP/), and the full description of the survey is provided in [Andrews et al. 2018](https://ui.adsabs.harvard.edu/abs/2018ApJ...869L..41A/abstract).
#
# ### Model Visibilities and MODEL_DATA
# In its simplest form, the measurement set just contains the visibility data in a ``DATA`` or ``CORRECTED_DATA`` column. In the process of using ``tclean`` to synthesize an image, however, CASA also calculates a set of model visibilities that correspond to the Fourier transform of the CLEAN model. It's possible to store these model visibilities to the measurement set if the ``tclean`` process was invoked with the ``savemodel="modelcolumn"`` parameter. The model visibilities will be stored in a ``MODEL_DATA`` column with the same shape as the ``DATA`` column.
#
# The calibrated DSHARP measurement sets available from the archive do not contain this ``MODEL_DATA`` column (most likely for space reasons), so we will need to recreate them by running the ``tclean`` algorithm with the relevant settings. The full reduction scripts are [available online](https://almascience.eso.org/almadata/lp/DSHARP/scripts/AS209_continuum.py), but we just need reproduce the relevant ``tclean`` commands used to produce a FITS image from the final, calibrated measurement set.
#
# Because the measurement set is large (0.9 Gb) and the ``tclean`` process is computationally expensive (taking about 1 hr on a single core), we have pre-executed those commands and cached the measurement set and ``tclean`` products into the ``AS209_MS`` local directory. If you're interested in the exact ``tclean`` commands used, please check out the [dl_and_tclean_AS209.py](dl_and_tclean_AS209.py) script directly.

import re
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import os

from traitlets.traitlets import validate

# change to the cached subdirectory that contains the cleaned MS
workdir = "AS209_MS"
os.chdir(workdir)
fname = "AS209_continuum.ms"
fitsname = "AS209.fits"

# ### Visualizing the CLEANed image
# Just to make sure the ``tclean`` process ran OK, let's check the synthesized image that was produced

from astropy.io import fits

hdul = fits.open(fitsname)
header = hdul[0].header
data = 1e3 * hdul[0].data  # mJy/pixel
# get the coordinate labels
nx = header["NAXIS1"]
ny = header["NAXIS2"]
# RA coordinates
CDELT1 = 3600 * header["CDELT1"]  # arcsec (converted from decimal deg)
CRPIX1 = header["CRPIX1"] - 1.0  # Now indexed from 0
# DEC coordinates
CDELT2 = 3600 * header["CDELT2"]  # arcsec
CRPIX2 = header["CRPIX2"] - 1.0  # Now indexed from 0
RA = (np.arange(nx) - nx / 2) * CDELT1  # [arcsec]
DEC = (np.arange(ny) - ny / 2) * CDELT2  # [arcsec]
# extent needs to include extra half-pixels.
# RA, DEC are pixel centers
ext = (
    RA[0] - CDELT1 / 2,
    RA[-1] + CDELT1 / 2,
    DEC[0] - CDELT2 / 2,
    DEC[-1] + CDELT2 / 2,
)  # [arcsec]
norm = matplotlib.colors.Normalize(vmin=0, vmax=np.max(data))

fig, ax = plt.subplots(nrows=1, figsize=(4.5, 3.5))
fig.subplots_adjust(left=0.2, bottom=0.2)
im = ax.imshow(data, extent=ext, origin="lower", animated=True, norm=norm)
cb = plt.colorbar(im, label="mJy / pixel")
r = 2.2
ax.set_xlim(r, -r)
ax.set_ylim(-r, r)
ax.set_xlabel(r"$\Delta \alpha \cos \delta$ [${}^{\prime\prime}$]")
ax.set_ylabel(r"$\Delta \delta$ [${}^{\prime\prime}$]")

# Great, it looks like things check out. Note that the main reason (at least for this tutorial) that we ran ``tclean`` was to generate the ``MODEL_DATA`` column in the measurement set. The actual CLEANed FITS image is just a nice byproduct.

# ## Examining measurement set structure
#
# Before you dive into the full analysis with CASA tools, it's a very good idea to inspect the measurement set using [listobs](https://casa.nrao.edu/casadocs-devel/stable/global-task-list/task_listobs/about).
#
# After you've done that, let's start exploring the visibility values. First we'll need to import and then instantiate the relevant CASA tools, [table](https://casa.nrao.edu/casadocs-devel/stable/global-tool-list/tool_table/methods) and [ms](https://casa.nrao.edu/casadocs-devel/stable/global-tool-list/tool_ms/methods).

import casatools

tb = casatools.table()
ms = casatools.ms()

# We can get the indexes of the unique spectral windows, which are typically indexed by the ``DATA_DESC_ID``.

tb.open(fname + "/DATA_DESCRIPTION")
SPECTRAL_WINDOW_ID = tb.getcol("SPECTRAL_WINDOW_ID")
tb.close()
print(SPECTRAL_WINDOW_ID)

# We see that there are 25 separate spectral windows! This is because the DSHARP images were produced using all available Band 6 continuum data on each source---not just the long baseline observations acquired in ALMA cycle 4. The merging of all of these individual observations is what creates this structure with so many spectral windows.

# Next, let's open the main table of the measurement set and inspect the column names

tb.open(fname)
colnames = tb.colnames()
tb.close()
print(colnames)

# Because there are multiple spectral windows which do not share the same dimensions, we cannot use the ``tb`` tool to read the data directly. If we try, we'll get an error.

try:
    tb.open(fname)
    weight = tb.getcol("WEIGHT")  # array of float64 with shape [npol, nvis]
    flag = tb.getcol("FLAG")  # array of bool with shape [npol, nchan, nvis]
    data = tb.getcol("DATA")  # array of complex128 with shape [npol, nchan, nvis]
except RuntimeError:
    print(
        "We can't use table tools here... the spws have different numbers of channels"
    )
finally:
    tb.close()

# So, we'll need to use the ``ms`` tool to read the visibilities for each spectral window, like so

ms.open(fname)
# select the spectral window
ms.selectinit(datadescid=0)
# query the desired columnames as a list
query = ms.getdata(["WEIGHT", "UVW", "DATA"])
# always a good idea to reset the earmarked data
ms.selectinit(reset=True)
ms.close()

# The returned query is a dictionary whose keys are the lowercase column names

print(query.keys())

# and whose values are the numerical arrays for the spectral window that we queried

print(query["data"])

# ## Using the tclean model to calculate residual visibilities
#
# In any data analysis where you're computing a forward model, it's a good consitency check to examine the data residuals from that model, and, in particular, whether their scatter matches the expectations from the noise properties.
#
# We can calculate data residuals using the model visibilities derived from the tclean model and stored in the ``MODEL_DATA`` column of the measurement set.

ms.open(fname)
# select the spectral window
ms.selectinit(datadescid=0)
# query the desired columnames as a list
query = ms.getdata(["MODEL_DATA"])
# always a good idea to reset the earmarked data
ms.selectinit(reset=True)
ms.close()

print(query["model_data"])

# Using these model visibilities, let's calculate the residuals for each polarization (XX, YY) in units of $\sigma$, where
#
# $$
# \sigma = \mathrm{sigma\_rescale} \times \sigma_0
# $$
#
# and
#
# $$
# \sigma_0 = \sqrt{1/w}
# $$
#
# The scatter is defined as
#
# $$
# \mathrm{scatter} = \frac{\mathrm{DATA} - \mathrm{MODEL\_DATA}}{\sigma}
# $$
#
# For now, $\mathrm{sigma\_rescale} = 1$, but we'll see why this parameter is needed in a moment.

# ### Helper functions for examining weight scatter
# Because we'd like to repeat this analysis for each spectral window in the measurement set, it makes things easier if we write these calculations as functions.
#
# The functions provided in this document are only dependent on the CASA tools ``tb`` and ``ms``. If you find yourself using these routines frequently, you might consider installing the *visread* package, since similar commands are provided in the API.

def get_scatter_datadescid(datadescid, sigma_rescale=1.0, apply_flags=True):
    """
    Calculate the scatter for each polarization.

    Args:
        datadescid (int): the DATA_DESC_ID to be queried
        sigma_rescale (int):  multiply the uncertainties by this factor
        apply_flags (bool): calculate the scatter *after* the flags have been applied

    Returns:
        scatter_XX, scatter_YY: a 2-tuple of numpy arrays containing the scatter in each polarization.
        If ``apply_flags==True``, each array will be 1-dimensional. If ``apply_flags==False``, each array
        will retain its original shape, including channelization (e.g., shape ``nchan,nvis``).
    """
    ms.open(fname)
    # select the key
    ms.selectinit(datadescid=datadescid)
    query = ms.getdata(
        ["DATA", "MODEL_DATA", "WEIGHT", "UVW", "ANTENNA1", "ANTENNA2", "FLAG"]
    )
    ms.selectinit(reset=True)
    ms.close()

    data, model_data, weight, flag = (
        query["data"],
        query["model_data"],
        query["weight"],
        query["flag"],
    )

    assert (
        len(model_data) > 0
    ), "MODEL_DATA column empty, retry tclean with savemodel='modelcolumn'"

    # subtract model from data
    residuals = data - model_data

    # calculate sigma from weight
    sigma = np.sqrt(1 / weight)
    sigma *= sigma_rescale

    # divide by weight, augmented for channel dim
    scatter = residuals / sigma[:, np.newaxis, :]

    # separate polarizations
    scatter_XX, scatter_YY = scatter
    flag_XX, flag_YY = flag

    if apply_flags:
        # flatten across channels
        scatter_XX = scatter_XX[~flag_XX]
        scatter_YY = scatter_YY[~flag_YY]

    return scatter_XX, scatter_YY

def gaussian(x, sigma=1):
    r"""
    Evaluate a reference Gaussian as a function of :math:`x`

    Args:
        x (float): location to evaluate Gaussian

    The Gaussian is defined as

    .. math::

        f(x) = \frac{1}{\sqrt{2 \pi}} \exp \left ( -\frac{x^2}{2}\right )

    Returns:
        Gaussian function evaluated at :math:`x`
    """
    return 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-0.5 * (x / sigma) ** 2)

def scatter_hist(scatter_XX, scatter_YY, log=False, **kwargs):
    """
    Plot a normalized histogram of scatter for real and imaginary
    components of XX and YY polarizations.

    Args:
        scatter_XX (1D numpy array)
        scatter_YY (1D numpy array)

    Returns:
        matplotlib figure
    """
    xs = np.linspace(-5, 5)

    figsize = kwargs.get("figsize", (6, 6))
    bins = kwargs.get("bins", 40)

    fig, ax = plt.subplots(ncols=2, nrows=2, figsize=figsize)
    ax[0, 0].hist(scatter_XX.real, bins=bins, density=True, log=log)
    ax[0, 0].set_xlabel(
        r"$\Re \{ V_\mathrm{XX} - \bar{V}_\mathrm{XX} \} / \sigma_\mathrm{XX}$"
    )
    ax[0, 1].hist(scatter_XX.imag, bins=bins, density=True, log=log)
    ax[0, 1].set_xlabel(
        r"$\Im \{ V_\mathrm{XX} - \bar{V}_\mathrm{XX} \} / \sigma_\mathrm{XX}$"
    )

    ax[1, 0].hist(scatter_YY.real, bins=bins, density=True, log=log)
    ax[1, 0].set_xlabel(
        r"$\Re \{ V_\mathrm{YY} - \bar{V}_\mathrm{YY} \} / \sigma_\mathrm{YY}$"
    )
    ax[1, 1].hist(scatter_YY.imag, bins=bins, density=True, log=log)
    ax[1, 1].set_xlabel(
        r"$\Im \{ V_\mathrm{YY} - \bar{V}_\mathrm{YY} \} / \sigma_\mathrm{YY}$"
    )

    for a in ax.flatten():
        a.plot(xs, gaussian(xs))

    fig.subplots_adjust(hspace=0.25, top=0.92)

    return fig

def plot_histogram_datadescid(
    datadescid, sigma_rescale=1.0, log=False, apply_flags=True
):
    """Wrap the scatter routine to plot a histogram of scatter for real and imaginary components of XX and YY polarizations, given a ``DATA_DESC_ID``.

    Args:
        datadescid (int): the DATA_DESC_ID to be queried
        sigma_rescale (int):  multiply the uncertainties by this factor
        log (bool): plot the histogram with a log stretch
        apply_flags (bool): calculate the scatter *after* the flags have been applied

    Returns:
        matplotlib figure
    """

    scatter_XX, scatter_YY = get_scatter_datadescid(
        datadescid=datadescid, sigma_rescale=sigma_rescale, apply_flags=apply_flags
    )

    scatter_XX = scatter_XX.flatten()
    scatter_YY = scatter_YY.flatten()

    fig = scatter_hist(scatter_XX, scatter_YY, log=log)
    fig.suptitle("DATA_DESC_ID: {:}".format(datadescid))

# ## Checking scatter for each spectral window
# Now lets use our helper functions to investigate the characteristics of each spectral window.

# ### Spectral Window 7: A correctly scaled SPW
# Let's see how the residual visibilities in spectral window 7 scatter relative to their expected Gaussian envelope

plot_histogram_datadescid(7, apply_flags=False)

# Great, it looks like things are pretty much as we would expect here.

# ### Spectral Window 22: Visibility outliers
# In the last example, we were a little bit cavalier and plotted the residuals for *all* visibilities, regardless of whether their ``FLAG`` was true. If we do the same for the visibilities in spectral window 22,

plot_histogram_datadescid(22, apply_flags=False)

# We find that something looks a little bit strange. Let's try plotting things on a log scale to get a closer look.

plot_histogram_datadescid(22, apply_flags=False, log=True)

# It appears as though there are a bunch of "outlier" visibilities included in this spectral window (the histogram y-axis is the relative frequency or probability density, so in actually there are a relatively small number of outlier visibilities). If the calibration and data preparation processes were done correctly, most likely, these visibilities are actually flagged. Let's try plotting only the valid, unflagged, visibilities

plot_histogram_datadescid(22, apply_flags=True, log=True)

# This is certainly an improvement, it looks like all of the egregious outlier visibilities were correctly flagged. However, we also see that the scatter in the visibility residuals is overdispersed relative to the envelope we would expect from the supplied weights. Plotting this on a normal scale might make the discrepancy more visible

plot_histogram_datadescid(22, apply_flags=True, log=False)

# Lets write a routine to estimate by how much the $\sigma_0$ values would need to be rescaled in order for the residual visibility scatter to match the expected reference Gaussian.

from scipy.optimize import minimize

def calculate_rescale_factor(scatter, **kwargs):
    bins = kwargs.get("bins", 40)
    bin_heights, bin_edges = np.histogram(scatter, density=True, bins=bins)
    bin_centers = bin_edges[:-1] + np.diff(bin_edges) / 2

    # find the sigma_rescale which minimizes the mean squared error
    # between the bin_heights and the expectations from the
    # reference Gaussian

    loss = lambda x: np.sum((bin_heights - gaussian(bin_centers, sigma=x)) ** 2)

    res = minimize(loss, 1.0)

    if res.success:
        return res.x[0]
    else:
        print(res)
        return False

def get_sigma_rescale_datadescid(datadescid, **kwargs):
    scatter_XX, scatter_YY = get_scatter_datadescid(
        datadescid, apply_flags=True, **kwargs
    )
    vals = np.array(
        [
            calculate_rescale_factor(scatter)
            for scatter in [
                scatter_XX.real,
                scatter_XX.imag,
                scatter_YY.real,
                scatter_YY.imag,
            ]
        ]
    )

    return np.average(vals)

factor = get_sigma_rescale_datadescid(22)
print(factor)

# If we rescale the $\sigma$ values to make them a factor of $\sim 1.85$ larger it looks like we are able to make the distributions match up a little bit better.

plot_histogram_datadescid(22, sigma_rescale=factor, apply_flags=True, log=False)

# It's not really clear why this factor works, but factors of $\sqrt{2}$ and 2 appear in changes to the CASA weight calculations ([which have changed across recent CASA versions](https://casa.nrao.edu/casadocs-devel/stable/calibration-and-visibility-data/data-weights)). Since this DSHARP measurement set contains visibilities acquired in previous ALMA cycles (and potentially calibrated with earlier versions of CASA), these changes in weight definitions may be responsible for the varying weight scatter across spectral window.
#
# To investigate the overdispersion in spectral window 22, lets plot up the rescaled visibilities in the $u,v$ plane, colorized by their residual values to see if we can discern any patterns that may be the result of an inadequate calibration or CLEAN model.

# get the baselines and flags for spectral window 22
ms.open(fname)
# select the key
ms.selectinit(datadescid=22)
query = ms.getdata(["UVW", "FLAG"])
ms.selectinit(reset=True)
ms.close()

flag_XX, flag_YY = query["flag"]
u, v, w = query["uvw"] * 1e-3  # [km]

# calculate the scatter of the residual visibilities
scatter_XX, scatter_YY = get_scatter_datadescid(
    22, sigma_rescale=factor, apply_flags=False
)

# Let's check the array shapes of each of these.

print(flag_XX.shape)
print(scatter_XX.shape)
print(u.shape)

# If we want to correctly apply the flags, we'll need to broadcast the baseline arrays to the full set of channels. (Even though this is a continuum dataset, it does have more than one channel to prevent bandwidth smearing).

nchan = flag_XX.shape[0]
broadcast = np.ones((nchan, 1))
uu = u * broadcast
vv = v * broadcast

# Now we index the "good" visibilities

uu_XX = uu[~flag_XX]
vv_XX = vv[~flag_YY]
scatter_XX = scatter_XX[~flag_XX]

uu_YY = uu[~flag_YY]
vv_YY = vv[~flag_YY]
scatter_YY = scatter_YY[~flag_YY]

vvmax = 5
norm = matplotlib.colors.Normalize(vmin=-vvmax, vmax=vvmax)

fig, ax = plt.subplots(nrows=1, figsize=(4, 4))
ax.scatter(uu_XX, vv_XX, s=0.1, c=scatter_XX.real, cmap="bwr", norm=norm)
im = ax.scatter(uu_YY, vv_YY, s=0.1, c=scatter_YY.real, cmap="bwr", norm=norm)
plt.colorbar(im, ax=ax)
ax.set_title("SPW: {:}".format(22))
ax.set_aspect("equal")
ax.set_xlabel(r"$u$ [km]")
ax.set_ylabel(r"$v$ [km]")

# We don't appear to see any large scale pattern with baseline, suggesting that the CLEAN model is doing a reasonable job of fitting the data across many spatial scales. We do see some of the largest outliers (most blue/red) at the beginning (or end) of the tracks. This probably has something to do with a less-than-optimal calibration at the beginning (or end) of the observation. Since there are no large systematic trends with baseline, we'll just accept these rescaled weights as is.

# ## Rescaling weights for export.
# We can use the previous routines to iterate through plots of each spectral window. For reference, here are the rescale factors we derived for each spectral window

sigma_rescale_factors = {
    ID: get_sigma_rescale_datadescid(ID) for ID in SPECTRAL_WINDOW_ID
}
for ID in SPECTRAL_WINDOW_ID:
    print(ID, "{:.2f}".format(sigma_rescale_factors[ID]))

# If you wanted to use these visibilities for forward modeling or [Regularized Maximum Likelihood (RML) imaging](https://mpol-dev.github.io/MPoL/), you will want to read and export the correctly flagged and rescaled visibilities and convert baselines to kilolambda.


"""
Tests for TimedeltaIndex methods behaving like their Timedelta counterparts
"""

import numpy as np
import pytest

from pandas._libs.tslibs.offsets import INVALID_FREQ_ERR_MSG

import pandas as pd
from pandas import Index, Series, Timedelta, TimedeltaIndex, timedelta_range
import pandas._testing as tm

class TestVectorizedTimedelta:
    def test_tdi_total_seconds(self):
        # GH#10939
        # test index
        rng = timedelta_range("1 days, 10:11:12.100123456", periods=2, freq="s")
        expt = [
            1 * 86400 + 10 * 3600 + 11 * 60 + 12 + 100123456.0 / 1e9,
            1 * 86400 + 10 * 3600 + 11 * 60 + 13 + 100123456.0 / 1e9,
        ]
        tm.assert_almost_equal(rng.total_seconds(), Index(expt))

        # test Series
        ser = Series(rng)
        s_expt = Series(expt, index=[0, 1])
        tm.assert_series_equal(ser.dt.total_seconds(), s_expt)

        # with nat
        ser[1] = np.nan
        s_expt = Series(
            [1 * 86400 + 10 * 3600 + 11 * 60 + 12 + 100123456.0 / 1e9, np.nan],
            index=[0, 1],
        )
        tm.assert_series_equal(ser.dt.total_seconds(), s_expt)

        # with both nat
        ser = Series([np.nan, np.nan], dtype="timedelta64[ns]")
        tm.assert_series_equal(
            ser.dt.total_seconds(), Series([np.nan, np.nan], index=[0, 1])
        )

    def test_tdi_round(self):
        td = pd.timedelta_range(start="16801 days", periods=5, freq="30Min")
        elt = td[1]

        expected_rng = TimedeltaIndex(
            [
                Timedelta("16801 days 00:00:00"),
                Timedelta("16801 days 00:00:00"),
                Timedelta("16801 days 01:00:00"),
                Timedelta("16801 days 02:00:00"),
                Timedelta("16801 days 02:00:00"),
            ]
        )
        expected_elt = expected_rng[1]

        tm.assert_index_equal(td.round(freq="H"), expected_rng)
        assert elt.round(freq="H") == expected_elt

        msg = INVALID_FREQ_ERR_MSG
        with pytest.raises(ValueError, match=msg):
            td.round(freq="foo")
        with pytest.raises(ValueError, match=msg):
            elt.round(freq="foo")

        msg = "<MonthEnd> is a non-fixed frequency"
        with pytest.raises(ValueError, match=msg):
            td.round(freq="M")
        with pytest.raises(ValueError, match=msg):
            elt.round(freq="M")

    @pytest.mark.parametrize(
        "freq,msg",
        [
            ("Y", "<YearEnd: month=12> is a non-fixed frequency"),
            ("M", "<MonthEnd> is a non-fixed frequency"),
            ("foobar", "Invalid frequency: foobar"),
        ],
    )
    def test_tdi_round_invalid(self, freq, msg):
        t1 = timedelta_range("1 days", periods=3, freq="1 min 2 s 3 us")

        with pytest.raises(ValueError, match=msg):
            t1.round(freq)
        with pytest.raises(ValueError, match=msg):
            # Same test for TimedeltaArray
            t1._data.round(freq)

    # TODO: de-duplicate with test_tdi_round
    def test_round(self):
        t1 = timedelta_range("1 days", periods=3, freq="1 min 2 s 3 us")
        t2 = -1 * t1
        t1a = timedelta_range("1 days", periods=3, freq="1 min 2 s")
        t1c = TimedeltaIndex([1, 1, 1], unit="D")

        # note that negative times round DOWN! so don't give whole numbers
        for (freq, s1, s2) in [
            ("N", t1, t2),
            ("U", t1, t2),
            (
                "L",
                t1a,
                TimedeltaIndex(
                    ["-1 days +00:00:00", "-2 days +23:58:58", "-2 days +23:57:56"]
                ),
            ),
            (
                "S",
                t1a,
                TimedeltaIndex(
                    ["-1 days +00:00:00", "-2 days +23:58:58", "-2 days +23:57:56"]
                ),
            ),
            ("12T", t1c, TimedeltaIndex(["-1 days", "-1 days", "-1 days"])),
            ("H", t1c, TimedeltaIndex(["-1 days", "-1 days", "-1 days"])),
            ("d", t1c, TimedeltaIndex([-1, -1, -1], unit="D")),
        ]:

            r1 = t1.round(freq)
            tm.assert_index_equal(r1, s1)
            r2 = t2.round(freq)
            tm.assert_index_equal(r2, s2)

    def test_components(self):
        rng = timedelta_range("1 days, 10:11:12", periods=2, freq="s")
        rng.components

        # with nat
        s = Series(rng)
        s[1] = np.nan

        result = s.dt.components
        assert not result.iloc[0].isna().all()
        assert result.iloc[1].isna().all()


from datetime import timedelta

import numpy as np
import pytest

import pandas as pd
from pandas import (
    DataFrame,
    Series,
)
import pandas._testing as tm
from pandas.core.indexes.timedeltas import timedelta_range

def test_asfreq_bug():
    df = DataFrame(data=[1, 3], index=[timedelta(), timedelta(minutes=3)])
    result = df.resample("1T").asfreq()
    expected = DataFrame(
        data=[1, np.nan, np.nan, 3],
        index=timedelta_range("0 day", periods=4, freq="1T"),
    )
    tm.assert_frame_equal(result, expected)

def test_resample_with_nat():
    # GH 13223
    index = pd.to_timedelta(["0s", pd.NaT, "2s"])
    result = DataFrame({"value": [2, 3, 5]}, index).resample("1s").mean()
    expected = DataFrame(
        {"value": [2.5, np.nan, 5.0]},
        index=timedelta_range("0 day", periods=3, freq="1S"),
    )
    tm.assert_frame_equal(result, expected)

def test_resample_as_freq_with_subperiod():
    # GH 13022
    index = timedelta_range("00:00:00", "00:10:00", freq="5T")
    df = DataFrame(data={"value": [1, 5, 10]}, index=index)
    result = df.resample("2T").asfreq()
    expected_data = {"value": [1, np.nan, np.nan, np.nan, np.nan, 10]}
    expected = DataFrame(
        data=expected_data, index=timedelta_range("00:00:00", "00:10:00", freq="2T")
    )
    tm.assert_frame_equal(result, expected)

def test_resample_with_timedeltas():

    expected = DataFrame({"A": np.arange(1480)})
    expected = expected.groupby(expected.index // 30).sum()
    expected.index = timedelta_range("0 days", freq="30T", periods=50)

    df = DataFrame(
        {"A": np.arange(1480)}, index=pd.to_timedelta(np.arange(1480), unit="T")
    )
    result = df.resample("30T").sum()

    tm.assert_frame_equal(result, expected)

    s = df["A"]
    result = s.resample("30T").sum()
    tm.assert_series_equal(result, expected["A"])

def test_resample_single_period_timedelta():

    s = Series(list(range(5)), index=timedelta_range("1 day", freq="s", periods=5))
    result = s.resample("2s").sum()
    expected = Series([1, 5, 4], index=timedelta_range("1 day", freq="2s", periods=3))
    tm.assert_series_equal(result, expected)

def test_resample_timedelta_idempotency():

    # GH 12072
    index = timedelta_range("0", periods=9, freq="10L")
    series = Series(range(9), index=index)
    result = series.resample("10L").mean()
    expected = series.astype(float)
    tm.assert_series_equal(result, expected)

def test_resample_offset_with_timedeltaindex():
    # GH 10530 & 31809
    rng = timedelta_range(start="0s", periods=25, freq="s")
    ts = Series(np.random.randn(len(rng)), index=rng)

    with_base = ts.resample("2s", offset="5s").mean()
    without_base = ts.resample("2s").mean()

    exp_without_base = timedelta_range(start="0s", end="25s", freq="2s")
    exp_with_base = timedelta_range(start="5s", end="29s", freq="2s")

    tm.assert_index_equal(without_base.index, exp_without_base)
    tm.assert_index_equal(with_base.index, exp_with_base)

def test_resample_categorical_data_with_timedeltaindex():
    # GH #12169
    df = DataFrame({"Group_obj": "A"}, index=pd.to_timedelta(list(range(20)), unit="s"))
    df["Group"] = df["Group_obj"].astype("category")
    result = df.resample("10s").agg(lambda x: (x.value_counts().index[0]))
    expected = DataFrame(
        {"Group_obj": ["A", "A"], "Group": ["A", "A"]},
        index=pd.TimedeltaIndex([0, 10], unit="s", freq="10s"),
    )
    expected = expected.reindex(["Group_obj", "Group"], axis=1)
    expected["Group"] = expected["Group_obj"]
    tm.assert_frame_equal(result, expected)

def test_resample_timedelta_values():
    # GH 13119
    # check that timedelta dtype is preserved when NaT values are
    # introduced by the resampling

    times = timedelta_range("1 day", "6 day", freq="4D")
    df = DataFrame({"time": times}, index=times)

    times2 = timedelta_range("1 day", "6 day", freq="2D")
    exp = Series(times2, index=times2, name="time")
    exp.iloc[1] = pd.NaT

    res = df.resample("2D").first()["time"]
    tm.assert_series_equal(res, exp)
    res = df["time"].resample("2D").first()
    tm.assert_series_equal(res, exp)

@pytest.mark.parametrize(
    "start, end, freq, resample_freq",
    [
        ("8H", "21h59min50s", "10S", "3H"),  # GH 30353 example
        ("3H", "22H", "1H", "5H"),
        ("527D", "5006D", "3D", "10D"),
        ("1D", "10D", "1D", "2D"),  # GH 13022 example
        # tests that worked before GH 33498:
        ("8H", "21h59min50s", "10S", "2H"),
        ("0H", "21h59min50s", "10S", "3H"),
        ("10D", "85D", "D", "2D"),
    ],
)
def test_resample_timedelta_edge_case(start, end, freq, resample_freq):
    # GH 33498
    # check that the timedelta bins does not contains an extra bin
    idx = timedelta_range(start=start, end=end, freq=freq)
    s = Series(np.arange(len(idx)), index=idx)
    result = s.resample(resample_freq).min()
    expected_index = timedelta_range(freq=resample_freq, start=start, end=end)
    tm.assert_index_equal(result.index, expected_index)
    assert result.index.freq == expected_index.freq
    assert not np.isnan(result[-1])

@pytest.mark.parametrize("duplicates", [True, False])
def test_resample_with_timedelta_yields_no_empty_groups(duplicates):
    # GH 10603
    df = DataFrame(
        np.random.normal(size=(10000, 4)),
        index=timedelta_range(start="0s", periods=10000, freq="3906250n"),
    )
    if duplicates:
        # case with non-unique columns
        df.columns = ["A", "B", "A", "C"]

    result = df.loc["1s":, :].resample("3s").apply(lambda x: len(x))

    expected = DataFrame(
        [[768] * 4] * 12 + [[528] * 4],
        index=timedelta_range(start="1s", periods=13, freq="3s"),
    )
    expected.columns = df.columns
    tm.assert_frame_equal(result, expected)

def test_resample_quantile_timedelta():
    # GH: 29485
    df = DataFrame(
        {"value": pd.to_timedelta(np.arange(4), unit="s")},
        index=pd.date_range("20200101", periods=4, tz="UTC"),
    )
    result = df.resample("2D").quantile(0.99)
    expected = DataFrame(
        {
            "value": [
                pd.Timedelta("0 days 00:00:00.990000"),
                pd.Timedelta("0 days 00:00:02.990000"),
            ]
        },
        index=pd.date_range("20200101", periods=2, tz="UTC", freq="2D"),
    )
    tm.assert_frame_equal(result, expected)