import os
import glob
import math
import random
import sys
import time
import numpy as np

from PIL import Image

from scipy.signal import savgol_filter
from six.moves import xrange

import umap
import argparse

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader, DistributedSampler

import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torchvision.utils import make_grid, save_image
import pytorch_lightning as pl

parser = argparse.ArgumentParser(description='VQ-VAE.')
parser.add_argument(
    '--rank',
    default=0,
    type=int,
    help="Rank of the training task"
)
parser.add_argument(
    '--world_size',
    default=1,
    type=int,
    help="World size of training tasks"
)
parser.add_argument(
    '--num_workers',
    default=4,
    type=int,
    help="Number of parallel data loaders"
)
parser.add_argument(
    '--init_method',
    default='tcp://192.168.1.154:23456',
    help="Master host (e.g. 'tcp://192.168.1.154:23456')"
)
parser.add_argument(
    '--epoch_start',
    default=1,
    type=int,
    help="Epoch to start training from (load savepoint)"
)
parser.add_argument(
    '--num_epochs',
    default=15000,
    type=int,
    help="Number of epochs to run"
)
parser.add_argument(
    '--batch_size',
    default=64,
    type=int,
    help="Number of data elements for one pass"
)
parser.add_argument(
    '--image_width',
    default=128,
    type=int,
    help="Horizontal image dimension"
)
parser.add_argument(
    '--image_height',
    default=128,
    type=int,
    help="Vertical image dimension"
)
parser.add_argument(
    '--backend',
    default="nccl",
    help="Distributed backend to use"
)

args = parser.parse_args()
print(args)

num_hiddens = 128
num_residual_hiddens = 32
num_residual_layers = 2
embedding_dim = 64
num_embeddings = 512
commitment_cost = 0.25
decay = 0.99
learning_rate = 1e-3
saved_models_path = "./saved_models/{epoch:08d}.vq-vae.net"
saved_results_path = "./results/{epoch:08d}.vq-vae.{name}"

if (torch.cuda.is_available()):
    device = "cuda"
    print("CUDA is available.")
else:
    device = "cpu"
    print("This example uses nccl as backend which is only available for (nVidia) GPUs.")
    sys.exit(-1)

pl.seed_everything(12345)

os.makedirs(os.path.dirname(saved_models_path.format(epoch=0)), exist_ok=True)
os.makedirs(os.path.dirname(saved_results_path.format(epoch=0,name='')), exist_ok=True)

class NoisySource_ImageDataset(Dataset):
    def __init__(self, images_pattern, transforms_before=None, transforms_after=None):
        self.images_pattern = images_pattern
        self.transforms_before = transforms_before
        self.transforms_after = transforms_after
        self.image_names = glob.glob(images_pattern)

    def __len__(self):
        return len(self.image_names)

    def __getitem__(self, idx):
        while True:
            img_name = self.image_names[idx]
            try:
                image_trg = self._load_image(img_name)
                if (image_trg.size[0] >= args.image_width) and (image_trg.size[1] >= args.image_height):
                    break
                print("Image {} too small ({}).".format(img_name, image_trg.size))
            except Exception as e:
                print("Exception {} - caught.".format(str(e)))
            idx = (idx+1) % len(self)

        if self.transforms_before:
            image_trg = self.transforms_before(image_trg)
        image_src = self._pixel_noise(image_trg)
        if self.transforms_after:
            image_trg = self.transforms_after(image_trg)
        if self.transforms_after:
            image_src = self.transforms_after(image_src)
        return {
            "image_src": image_src,
            "image_trg": image_trg,
        }

    def _load_image(self, path):
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('RGB')

    def _pixel_noise(self, src):
        factor = random.randint(8,64)
        pix = src.copy()
        nsiz_w = max(int(src.width/factor),1)
        nsiz_h = max(int(src.height/factor),1)
        pix = pix.resize((nsiz_w, nsiz_h))
        pix = pix.resize(src.size, Image.NEAREST)
        w0 = random.randint(0,src.width)
        h0 = random.randint(0,src.height)
        w1 = random.randint(w0, src.width)
        h1 = random.randint(h0, src.height)
        pix = pix.crop((w0,h0,w1,h1))
        res = src.copy()
        res.paste(pix, (w0,h0))
        return res

# Training Data
# Dataset

training_data_dataset = NoisySource_ImageDataset(
    "/data/imagenet/imagenet_images/*/*.jpg",
    transforms_before=transforms.Compose(
        [
            transforms.RandomCrop((args.image_height,args.image_width), padding_mode='reflect')
        ]
    ),
    transforms_after=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
        ]
    ),
)

# training_data_dataset = datasets.ImageFolder(
#     "/data/imagenet",
#     transform=transforms.Compose(
#         [
#             transforms.Resize((args.image_height,args.image_width), interpolation=2),
#             transforms.ToTensor(),
#             transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
#         ]
#     ),
# #    target_transform=None,
# #    loader=<function default_loader>,
# #    is_valid_file=None
# )
# training_data_dataset = datasets.CIFAR10(
#     root="data", train=True, download=True,
#     transform=transforms.Compose(
#         [
#             transforms.Resize(IMAGE_SIZE, interpolation=2),
#             transforms.ToTensor(),
#             transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
#         ]
#     )
# )
# training_data_dataset = datasets.MNIST(
#     root="data",
#     train=True,
#     transform=transforms.Compose(
#         [
#             transforms.Resize(IMAGE_SIZE, interpolation=2),
#             transforms.Grayscale(num_output_channels=3),
#             transforms.ToTensor(),
#             transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
#         ]
#     ),
#     target_transform=None,
#     download=True
# )
# DataSampler
training_data_sampler = DistributedSampler(
    training_data_dataset,
    num_replicas=args.world_size,
    rank=args.rank,
    shuffle=True
)
# DataLoader
training_data_loader = DataLoader(
    training_data_dataset,
    batch_size=args.batch_size,
    shuffle=(training_data_sampler is None),
    sampler=training_data_sampler,
    pin_memory=True,
    num_workers=args.num_workers,
)

# Validation Data
# Dataset
validation_data_dataset = NoisySource_ImageDataset(
    "/data/imagenet/imagenet_images/*/*.jpg",
    transforms_before=transforms.Compose(
        [
            transforms.RandomCrop((args.image_height,args.image_width), )
        ]
    ),
    transforms_after=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
        ]
    ),
)
# validation_data_dataset = datasets.ImageFolder(
#     "/data/imagenet",
#     transform=transforms.Compose(
#         [
#             transforms.Resize((args.image_height,args.image_width), interpolation=2),
#             transforms.ToTensor(),
#             transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
#         ]
#     ),
# #    target_transform=None,
# #    loader=<function default_loader>,
# #    is_valid_file=None
# )
# validation_data_dataset = datasets.CIFAR10(
#     root="data", train=False, download=True,
#     transform=transforms.Compose(
#         [
#             transforms.Resize(IMAGE_SIZE, interpolation=2),
#             transforms.ToTensor(),
#             transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
#         ]
#     )
# )
# validation_data_dataset = datasets.MNIST(
#     root="data",
#     train=False,
#     transform=transforms.Compose(
#         [
#             transforms.Resize(IMAGE_SIZE, interpolation=2),
#             transforms.Grayscale(num_output_channels=3),
#             transforms.ToTensor(),
#             transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
#         ]
#     ),
#     target_transform=None,
#     download=True
# )
# DataSampler
validation_data_sampler = DistributedSampler(
    validation_data_dataset,
    num_replicas=args.world_size,
    rank=args.rank,
    shuffle=True
)
# DataLoader
validation_data_loader = DataLoader(
    validation_data_dataset,
    batch_size=32,
    shuffle=(validation_data_sampler is None),
    sampler=validation_data_sampler,
    pin_memory=True,
)

training_data_loader_iterator = iter(training_data_loader)
validation_data_loader_iterator = iter(validation_data_loader)

# Vector Quantizer Layer
#
# This layer takes a tensor to be quantized. The channel dimension will be used as the space in which to quantize. All other dimensions will be flattened and will be seen as different examples to quantize.
#
# The output tensor will have the same shape as the input.
#
# As an example for a `BCHW` tensor of shape `[16, 64, 32, 32]`, we will first convert it to an `BHWC` tensor of shape `[16, 32, 32, 64]` and then reshape it into `[16384, 64]` and all `16384` vectors of size `64`  will be quantized independently. In otherwords, the channels are used as the space in which to quantize. All other dimensions will be flattened and be seen as different examples to quantize, `16384` in this case.

class VectorQuantizer(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, commitment_cost):
        super(VectorQuantizer, self).__init__()

        self._embedding_dim = embedding_dim
        self._num_embeddings = num_embeddings

        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
        self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings)
        self._commitment_cost = commitment_cost

    def forward(self, inputs):
        # convert inputs from BCHW -> BHWC
        inputs = inputs.permute(0, 2, 3, 1).contiguous()
        input_shape = inputs.shape

        # Flatten input
        flat_input = inputs.view(-1, self._embedding_dim)

        # Calculate distances
        distances = (torch.sum(flat_input**2, dim=1, keepdim=True)
                    + torch.sum(self._embedding.weight**2, dim=1)
                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))

        # Encoding
        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
        encodings.scatter_(1, encoding_indices, 1)

        # Quantize and unflatten
        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)

        # Loss
        e_latent_loss = F.mse_loss(quantized.detach(), inputs)
        q_latent_loss = F.mse_loss(quantized, inputs.detach())
        loss = q_latent_loss + self._commitment_cost * e_latent_loss

        quantized = inputs + (quantized - inputs).detach()
        avg_probs = torch.mean(encodings, dim=0)
        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))

        # convert quantized from BHWC -> BCHW
        return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings

# We will also implement a slightly modified version  which will use exponential moving averages to update the embedding vectors instead of an auxillary loss. This has the advantage that the embedding updates are independent of the choice of optimizer for the encoder, decoder and other parts of the architecture. For most experiments the EMA version trains faster than the non-EMA version.

class VectorQuantizerEMA(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5):
        super(VectorQuantizerEMA, self).__init__()

        self._embedding_dim = embedding_dim
        self._num_embeddings = num_embeddings

        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
        self._embedding.weight.data.normal_()
        self._commitment_cost = commitment_cost

        self.register_buffer('_ema_cluster_size', torch.zeros(num_embeddings))
        self._ema_w = nn.Parameter(torch.Tensor(num_embeddings, self._embedding_dim))
        self._ema_w.data.normal_()

        self._decay = decay
        self._epsilon = epsilon

    def forward(self, inputs):
        # convert inputs from BCHW -> BHWC
        inputs = inputs.permute(0, 2, 3, 1).contiguous()
        input_shape = inputs.shape

        # Flatten input
        flat_input = inputs.view(-1, self._embedding_dim)

        # Calculate distances
        distances = (torch.sum(flat_input**2, dim=1, keepdim=True)
                    + torch.sum(self._embedding.weight**2, dim=1)
                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))

        # Encoding
        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
        encodings.scatter_(1, encoding_indices, 1)

        # Quantize and unflatten
        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)

        # Use EMA to update the embedding vectors
        if self.training:
            self._ema_cluster_size = self._ema_cluster_size * self._decay +                                      (1 - self._decay) * torch.sum(encodings, 0)

            # Laplace smoothing of the cluster size
            n = torch.sum(self._ema_cluster_size.data)
            self._ema_cluster_size = (
                (self._ema_cluster_size + self._epsilon)
                / (n + self._num_embeddings * self._epsilon) * n)

            dw = torch.matmul(encodings.t(), flat_input)
            self._ema_w = nn.Parameter(self._ema_w * self._decay + (1 - self._decay) * dw)

            self._embedding.weight = nn.Parameter(self._ema_w / self._ema_cluster_size.unsqueeze(1))

        # Loss
        e_latent_loss = F.mse_loss(quantized.detach(), inputs)
        loss = self._commitment_cost * e_latent_loss

        # Straight Through Estimator
        quantized = inputs + (quantized - inputs).detach()
        avg_probs = torch.mean(encodings, dim=0)
        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))

        # convert quantized from BHWC -> BCHW
        return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings

# Encoder & Decoder Architecture
#
# The encoder and decoder architecture is based on a ResNet and is implemented below:

class Residual(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_hiddens):
        super(Residual, self).__init__()
        self._block = nn.Sequential(
            nn.ReLU(True),
            nn.Conv2d(in_channels=in_channels,
                      out_channels=num_residual_hiddens,
                      kernel_size=3, stride=1, padding=1, bias=False),
            nn.ReLU(True),
            nn.Conv2d(in_channels=num_residual_hiddens,
                      out_channels=num_hiddens,
                      kernel_size=1, stride=1, bias=False)
        )

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

class ResidualStack(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
        super(ResidualStack, self).__init__()
        self._num_residual_layers = num_residual_layers
        self._layers = nn.ModuleList([Residual(in_channels, num_hiddens, num_residual_hiddens)
                             for _ in range(self._num_residual_layers)])

    def forward(self, x):
        for i in range(self._num_residual_layers):
            x = self._layers[i](x)
        return F.relu(x)

class Encoder(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
        super(Encoder, self).__init__()

        self._conv_1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=num_hiddens//2,
            kernel_size=4,
            stride=2, padding=1
        )
        self._conv_2 = nn.Conv2d(
            in_channels=num_hiddens//2,
            out_channels=num_hiddens,
            kernel_size=4,
            stride=2, padding=1
        )
        self._conv_3 = nn.Conv2d(
            in_channels=num_hiddens,
            out_channels=num_hiddens,
            kernel_size=3,
            stride=1, padding=1
        )
        self._residual_stack = ResidualStack(
            in_channels=num_hiddens,
            num_hiddens=num_hiddens,
            num_residual_layers=num_residual_layers,
            num_residual_hiddens=num_residual_hiddens
        )

    def forward(self, inputs):
        x = self._conv_1(inputs)
        x = F.relu(x)

        x = self._conv_2(x)
        x = F.relu(x)

        x = self._conv_3(x)
        return self._residual_stack(x)

class Decoder(nn.Module):
    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
        super(Decoder, self).__init__()

        self._conv_1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=num_hiddens,
            kernel_size=3,
            stride=1, padding=1
        )

        self._residual_stack = ResidualStack(
            in_channels=num_hiddens,
            num_hiddens=num_hiddens,
            num_residual_layers=num_residual_layers,
            num_residual_hiddens=num_residual_hiddens
        )

        self._conv_trans_1 = nn.ConvTranspose2d(
            in_channels=num_hiddens,
            out_channels=num_hiddens//2,
            kernel_size=4,
            stride=2, padding=1
        )

        self._conv_trans_2 = nn.ConvTranspose2d(
            in_channels=num_hiddens//2,
            out_channels=3,
            kernel_size=4,
            stride=2, padding=1
        )

    def forward(self, inputs):
        x = self._conv_1(inputs)

        x = self._residual_stack(x)

        x = self._conv_trans_1(x)
        x = F.relu(x)

        return self._conv_trans_2(x)

# Model & Optimizer
class Model(nn.Module):
    def __init__(self, num_hiddens, num_residual_layers, num_residual_hiddens,
                 num_embeddings, embedding_dim, commitment_cost, decay=0):
        super(Model, self).__init__()

        self._encoder = Encoder(
            3, num_hiddens,
            num_residual_layers,
            num_residual_hiddens
        )
        self._pre_vq_conv = nn.Conv2d(
            in_channels=num_hiddens,
            out_channels=embedding_dim,
            kernel_size=1,
            stride=1
        )
        if decay > 0.0:
            self._vq_vae = VectorQuantizerEMA(
                num_embeddings, embedding_dim,
                commitment_cost, decay
            )
        else:
            self._vq_vae = VectorQuantizer(
                num_embeddings, embedding_dim,
                commitment_cost
            )
        self._decoder = Decoder(
            embedding_dim,
            num_hiddens,
            num_residual_layers,
            num_residual_hiddens
        )

    def forward(self, x):
        z = self._encoder(x)
        z = self._pre_vq_conv(z)
        loss, quantized, perplexity, _ = self._vq_vae(z)
        x_recon = self._decoder(quantized)
        return loss, x_recon, perplexity

model = Model(
    num_hiddens, num_residual_layers, num_residual_hiddens,
    num_embeddings, embedding_dim,
    commitment_cost, decay
).to(device)

load_model_path = saved_models_path.format(epoch=args.epoch_start)
if os.path.exists(load_model_path):
    print("Loading {}".format(load_model_path))
    load_dict = torch.load(load_model_path)
    model.load_state_dict(load_dict['model_state_dict'])
    print(load_dict.keys())
else:
    print("Could not read {}; no data was loaded.".format(load_model_path))

optimizer = optim.Adam(model.parameters(), lr=learning_rate, amsgrad=False)

def save_model(model, epoch):
    save_dict = {
        'model_state_dict': model.state_dict(),
        'epoch': epoch,
        'args': repr(args),
    }
    torch.save(save_dict, saved_models_path.format(epoch=epoch))

# Show reconstructions
def validate():
    global validation_data_loader_iterator
    model.eval()

    # (valid_originals, _) = next(iter(validation_data_loader))
    while True:
        try:
            data = next(validation_data_loader_iterator)
            break
        except StopIteration:
            validation_data_loader_iterator = iter(validation_data_loader)
    valid_originals = data['image_src']
    valid_originals_targets = data['image_trg']
    valid_originals = valid_originals.to(device)
    vq_output_eval = model._pre_vq_conv(model._encoder(valid_originals))
    _, valid_quantize, _, _ = model._vq_vae(vq_output_eval)
    valid_reconstructions = model._decoder(valid_quantize)

    save_image(
        valid_reconstructions.cpu().data+0.5,
        fp=saved_results_path.format(epoch=epoch, name='reconstruction.png'),
        nrow=8,
        padding=2,
        normalize=False,
        range=None,
        scale_each=False,
        pad_value=0,
        format="png"
    )
    save_image(
        valid_originals.cpu()+0.5,
        fp=saved_results_path.format(epoch=epoch, name='originals.png'),
        nrow=8,
        padding=2,
        normalize=False,
        range=None,
        scale_each=False,
        pad_value=0,
        format="png"
    )
    save_image(
        valid_originals_targets.cpu()+0.5,
        fp=saved_results_path.format(epoch=epoch, name='targets.png'),
        nrow=8,
        padding=2,
        normalize=False,
        range=None,
        scale_each=False,
        pad_value=0,
        format="png"
    )
    model.train()

def convert_size(size_bytes):
   if size_bytes == 0:
       return "0B"
   size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
   i = int(math.floor(math.log(size_bytes, 1024)))
   p = math.pow(1024, i)
   s = round(size_bytes / p, 2)
   return "%s %s" % (s, size_name[i])

def device_memory_info(device_number=0):
    t = torch.cuda.get_device_properties(device_number).total_memory
    c = torch.cuda.memory_cached(device_number)
    a = torch.cuda.memory_allocated(device_number)
    f = c-a  # free inside cache
    return {
        'total': t,
        'cached': c,
        'allocated': a,
        'free': f,
    }

def print_device_memory_info():
    info = device_memory_info()
    print(
        "Total: {}, Cached: {}, Allocated: {}, Free: {}".format(
            convert_size(info['total']),
            convert_size(info['cached']),
            convert_size(info['allocated']),
            convert_size(info['free']),
        )
    )

class Timer(object):
    def __init__(self):
        self.time_instantiated = time.time()
        self.time_last = None
        self.kwargs_last = None

    def start(self, **kwargs):
        self.kwargs_last = kwargs
        self.time_last = time.time()
        return {
            'span': 0,
            'kwargs': kwargs,
        }

    def lap(self, **kwargs):
        time_now = time.time()
        span = time_now - self.time_last
        kwargs = self.kwargs_last
        self.time_last = time_now
        self.kwargs_last = kwargs
        result = kwargs
        result['span'] = span
        return result

# Train

print_device_memory_info()

print("Initializing process group...")
dist.init_process_group(
    backend=args.backend,
    init_method=args.init_method,
    rank=args.rank,
    world_size=args.world_size
)

print("Trainig starts!")
model.train()
train_res_recon_error = []
train_res_perplexity = []

timer = Timer()
timer.start(iterations=args.epoch_start)

for epoch in xrange(args.epoch_start, args.epoch_start+args.num_epochs+1):
    try:
        data = next(training_data_loader_iterator)
    except StopIteration:
        training_data_loader_iterator = iter(training_data_loader)
        data = next(training_data_loader_iterator)
    image_src = data['image_src']
    image_trg = data['image_trg']
    image_src = data['image_src'].to(device)
    image_trg = data['image_trg'].to(device)
    optimizer.zero_grad()

    vq_loss, data_recon, perplexity = model(image_src)
    recon_error = F.mse_loss(data_recon, image_trg)
    loss = recon_error + vq_loss
    loss = loss / float(len(data['image_src']))
    loss.backward()
    optimizer.step()

    train_res_recon_error.append(recon_error.item())
    train_res_perplexity.append(perplexity.item())

    for param in model.parameters():
        if param.grad is not None:
            dist.all_reduce(param.grad.data, op=torch.distributed.ReduceOp.SUM)
            param.grad.data /= float(args.world_size)

    if ((epoch % 100 == 0) or (epoch==args.epoch_start+args.num_epochs)) and (epoch != args.epoch_start):
        timer_info = timer.lap(iterations=epoch)
        fract = len(image_trg) * (epoch - timer_info['iterations']) / timer_info['span']

        # from IPython.display import clear_output
        # clear_output(wait=True)
        print(
            '{iterations:d} iterations, {fract:f}/s'.format(
                iterations=epoch,
                fract=fract
            )
        )
        print('recon_error: %.3f' % np.mean(train_res_recon_error[-100:]))
        print('perplexity: %.3f' % np.mean(train_res_perplexity[-100:]))
        save_model(model, epoch)
        validate()

        print_device_memory_info()
        timer.start(iterations=epoch)

print("Done.")

# # Plot Loss
# train_res_recon_error_smooth = savgol_filter(train_res_recon_error, 201, 7)
# train_res_perplexity_smooth = savgol_filter(train_res_perplexity, 201, 7)
#
# f = plt.figure(figsize=(16,8))
# ax = f.add_subplot(1,2,1)
# ax.plot(train_res_recon_error_smooth)
# ax.set_yscale('log')
# ax.set_title('Smoothed NMSE.')
# ax.set_xlabel('iteration')
#
# ax = f.add_subplot(1,2,2)
# ax.plot(train_res_perplexity_smooth)
# ax.set_title('Smoothed Average codebook usage (perplexity).')
# ax.set_xlabel('iteration')

# def show(img):
#     npimg = img.numpy()
#     fig = plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')
#     fig.axes.get_xaxis().set_visible(False)
#     fig.axes.get_yaxis().set_visible(False)

# show(make_grid(valid_reconstructions.cpu().data+0.5, range=(0.0, 1.0), scale_each=True))
# show(make_grid(valid_originals.cpu()+0.5))

# # View Embedding
# proj = umap.UMAP(n_neighbors=3,
#                  min_dist=0.1,
#                  metric='cosine').fit_transform(model._vq_vae._embedding.weight.data.cpu())
#
# plt.scatter(proj[:,0], proj[:,1], alpha=0.3)
#

# python vq-vae.py --rank 0 --world_size 2 --epoch_start 0 --num_epochs 15000 --init_method 'tcp://192.168.1.154:23456'
# python vq-vae.py --rank 1 --world_size 2 --epoch_start 0 --num_epochs 15000 --init_method 'tcp://192.168.1.154:23456'


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Module which provides classes to perform PSF Photometry"""

import numpy as np
import warnings

from astropy.modeling.fitting import LevMarLSQFitter
from astropy.nddata.utils import overlap_slices
from astropy.stats import gaussian_sigma_to_fwhm, SigmaClip
from astropy.table import Table, Column, vstack, hstack
from astropy.utils.exceptions import AstropyUserWarning

from .groupstars import DAOGroup
from .utils import (get_grouped_psf_model, subtract_psf,
                    _extract_psf_fitting_names)
from ..aperture import CircularAperture, aperture_photometry
from ..background import MMMBackground
from ..detection import DAOStarFinder

__all__ = ['BasicPSFPhotometry', 'IterativelySubtractedPSFPhotometry',
           'DAOPhotPSFPhotometry']

class BasicPSFPhotometry:
    """
    This class implements a PSF photometry algorithm that can find
    sources in an image, group overlapping sources into a single model,
    fit the model to the sources, and subtracting the models from the
    image. This is roughly equivalent to the DAOPHOT routines FIND,
    GROUP, NSTAR, and SUBTRACT.  This implementation allows a flexible
    and customizable interface to perform photometry. For instance, one
    is able to use different implementations for grouping and finding
    sources by using ``group_maker`` and ``finder`` respectivelly. In
    addition, sky background estimation is performed by
    ``bkg_estimator``.

    Parameters
    ----------
    group_maker : callable or `~photutils.psf.GroupStarsBase`
        ``group_maker`` should be able to decide whether a given star
        overlaps with any other and label them as beloging to the same
        group.  ``group_maker`` receives as input an
        `~astropy.table.Table` object with columns named as ``id``,
        ``x_0``, ``y_0``, in which ``x_0`` and ``y_0`` have the same
        meaning of ``xcentroid`` and ``ycentroid``.  This callable must
        return an `~astropy.table.Table` with columns ``id``, ``x_0``,
        ``y_0``, and ``group_id``. The column ``group_id`` should cotain
        integers starting from ``1`` that indicate which group a given
        source belongs to. See, e.g., `~photutils.psf.DAOGroup`.
    bkg_estimator : callable, instance of any `~photutils.BackgroundBase` subclass, or None
        ``bkg_estimator`` should be able to compute either a scalar
        background or a 2D background of a given 2D image. See, e.g.,
        `~photutils.background.MedianBackground`.  If None, no
        background subtraction is performed.
    psf_model : `astropy.modeling.Fittable2DModel` instance
        PSF or PRF model to fit the data. Could be one of the models in
        this package like `~photutils.psf.sandbox.DiscretePRF`,
        `~photutils.psf.IntegratedGaussianPRF`, or any other suitable 2D
        model.  This object needs to identify three parameters (position
        of center in x and y coordinates and the flux) in order to set
        them to suitable starting values for each fit. The names of
        these parameters should be given as ``x_0``, ``y_0`` and
        ``flux``.  `~photutils.psf.prepare_psf_model` can be used to
        prepare any 2D model to match this assumption.
    fitshape : int or length-2 array-like
        Rectangular shape around the center of a star which will be used
        to collect the data to do the fitting. Can be an integer to be
        the same along both axes. E.g., 5 is the same as (5, 5), which
        means to fit only at the following relative pixel positions:
        [-2, -1, 0, 1, 2].  Each element of ``fitshape`` must be an odd
        number.
    finder : callable or instance of any `~photutils.detection.StarFinderBase` subclasses or None
        ``finder`` should be able to identify stars, i.e. compute a
        rough estimate of the centroids, in a given 2D image.
        ``finder`` receives as input a 2D image and returns an
        `~astropy.table.Table` object which contains columns with names:
        ``id``, ``xcentroid``, ``ycentroid``, and ``flux``. In which
        ``id`` is an integer-valued column starting from ``1``,
        ``xcentroid`` and ``ycentroid`` are center position estimates of
        the sources and ``flux`` contains flux estimates of the sources.
        See, e.g., `~photutils.detection.DAOStarFinder`.  If ``finder``
        is ``None``, initial guesses for positions of objects must be
        provided.
    fitter : `~astropy.modeling.fitting.Fitter` instance
        Fitter object used to compute the optimized centroid positions
        and/or flux of the identified sources. See
        `~astropy.modeling.fitting` for more details on fitters.
    aperture_radius : float or None
        The radius (in units of pixels) used to compute initial
        estimates for the fluxes of sources. If ``None``, one FWHM will
        be used if it can be determined from the ``psf_model``.

    Notes
    -----
    Note that an ambiguity arises whenever ``finder`` and
    ``init_guesses`` (keyword argument for ``do_photometry``) are both
    not ``None``. In this case, ``finder`` is ignored and initial
    guesses are taken from ``init_guesses``. In addition, an warning is
    raised to remaind the user about this behavior.

    If there are problems with fitting large groups, change the
    parameters of the grouping algorithm to reduce the number of sources
    in each group or input a ``star_groups`` table that only includes
    the groups that are relevant (e.g. manually remove all entries that
    coincide with artifacts).

    References
    ----------
    [1] Stetson, Astronomical Society of the Pacific, Publications,
        (ISSN 0004-6280), vol. 99, March 1987, p. 191-222.
        Available at: http://adsabs.harvard.edu/abs/1987PASP...99..191S
    """

    def __init__(self, group_maker, bkg_estimator, psf_model, fitshape,
                 finder=None, fitter=LevMarLSQFitter(), aperture_radius=None):
        self.group_maker = group_maker
        self.bkg_estimator = bkg_estimator
        self.psf_model = psf_model
        self.fitter = fitter
        self.fitshape = fitshape
        self.finder = finder
        self.aperture_radius = aperture_radius
        self._pars_to_set = None
        self._pars_to_output = None
        self._residual_image = None

    @property
    def fitshape(self):
        return self._fitshape

    @fitshape.setter
    def fitshape(self, value):
        value = np.asarray(value)

        # assume a lone value should mean both axes
        if value.shape == ():
            value = np.array((value, value))

        if value.size == 2:
            if np.all(value) > 0:
                if np.all(value % 2) == 1:
                    self._fitshape = tuple(value)
                else:
                    raise ValueError('fitshape must be odd integer-valued, '
                                     'received fitshape = {}'.format(value))
            else:
                raise ValueError('fitshape must have positive elements, '
                                 'received fitshape = {}'.format(value))
        else:
            raise ValueError('fitshape must have two dimensions, '
                             'received fitshape = {}'.format(value))

    @property
    def aperture_radius(self):
        return self._aperture_radius

    @aperture_radius.setter
    def aperture_radius(self, value):
        if isinstance(value, (int, float)) and value > 0:
            self._aperture_radius = value
        elif value is None:
            self._aperture_radius = value
        else:
            raise ValueError('aperture_radius must be a real-valued '
                             'number, received aperture_radius = {}'
                             .format(value))

    def get_residual_image(self):
        """
        Returns an image that is the result of the subtraction between
        the original image and the fitted sources.

        Returns
        -------
        residual_image : 2D array-like, `~astropy.io.fits.ImageHDU`, `~astropy.io.fits.HDUList`
        """

        return self._residual_image

    def __call__(self, image, init_guesses=None):
        """
        Performs PSF photometry. See `do_photometry` for more details
        including the `__call__` signature.
        """

        return self.do_photometry(image, init_guesses)

    def do_photometry(self, image, init_guesses=None):
        """
        Perform PSF photometry in ``image``.

        This method assumes that ``psf_model`` has centroids and flux
        parameters which will be fitted to the data provided in
        ``image``. A compound model, in fact a sum of ``psf_model``,
        will be fitted to groups of stars automatically identified by
        ``group_maker``. Also, ``image`` is not assumed to be background
        subtracted.  If ``init_guesses`` are not ``None`` then this
        method uses ``init_guesses`` as initial guesses for the
        centroids. If the centroid positions are set as ``fixed`` in the
        PSF model ``psf_model``, then the optimizer will only consider
        the flux as a variable.

        Parameters
        ----------
        image : 2D array-like, `~astropy.io.fits.ImageHDU`, `~astropy.io.fits.HDUList`
            Image to perform photometry.
        init_guesses: `~astropy.table.Table`
            Table which contains the initial guesses (estimates) for the
            set of parameters. Columns 'x_0' and 'y_0' which represent
            the positions (in pixel coordinates) for each object must be
            present.  'flux_0' can also be provided to set initial
            fluxes.  If 'flux_0' is not provided, aperture photometry is
            used to estimate initial values for the fluxes. Additional
            columns of the form '<parametername>_0' will be used to set
            the initial guess for any parameters of the ``psf_model``
            model that are not fixed.

        Returns
        -------
        output_tab : `~astropy.table.Table` or None
            Table with the photometry results, i.e., centroids and
            fluxes estimations and the initial estimates used to start
            the fitting process. Uncertainties on the fitted parameters
            are reported as columns called ``<paramname>_unc`` provided
            that the fitter object contains a dictionary called
            ``fit_info`` with the key ``param_cov``, which contains the
            covariance matrix. If ``param_cov`` is not present,
            uncertanties are not reported.
        """

        if self.bkg_estimator is not None:
            image = image - self.bkg_estimator(image)

        if self.aperture_radius is None:
            if hasattr(self.psf_model, 'fwhm'):
                self.aperture_radius = self.psf_model.fwhm.value
            elif hasattr(self.psf_model, 'sigma'):
                self.aperture_radius = (self.psf_model.sigma.value *
                                        gaussian_sigma_to_fwhm)

        if init_guesses is not None:
            # make sure the code does not modify user's input
            init_guesses = init_guesses.copy()
            if self.aperture_radius is None:
                if 'flux_0' not in init_guesses.colnames:
                    raise ValueError('aperture_radius is None and could not '
                                     'be determined by psf_model. Please, '
                                     'either provided a value for '
                                     'aperture_radius or define fwhm/sigma '
                                     'at psf_model.')

            if self.finder is not None:
                warnings.warn('Both init_guesses and finder are different '
                              'than None, which is ambiguous. finder is '
                              'going to be ignored.', AstropyUserWarning)

            if 'flux_0' not in init_guesses.colnames:
                apertures = CircularAperture((init_guesses['x_0'],
                                              init_guesses['y_0']),
                                             r=self.aperture_radius)

                init_guesses['flux_0'] = aperture_photometry(
                    image, apertures)['aperture_sum']
        else:
            if self.finder is None:
                raise ValueError('Finder cannot be None if init_guesses are '
                                 'not given.')
            sources = self.finder(image)
            if len(sources) > 0:
                apertures = CircularAperture((sources['xcentroid'],
                                              sources['ycentroid']),
                                             r=self.aperture_radius)

                sources['aperture_flux'] = aperture_photometry(
                    image, apertures)['aperture_sum']

                init_guesses = Table(names=['x_0', 'y_0', 'flux_0'],
                                     data=[sources['xcentroid'],
                                           sources['ycentroid'],
                                           sources['aperture_flux']])

        self._define_fit_param_names()
        for p0, param in self._pars_to_set.items():
            if p0 not in init_guesses.colnames:
                init_guesses[p0] = (len(init_guesses) *
                                    [getattr(self.psf_model, param).value])

        star_groups = self.group_maker(init_guesses)
        output_tab, self._residual_image = self.nstar(image, star_groups)

        star_groups = star_groups.group_by('group_id')
        output_tab = hstack([star_groups, output_tab])

        return output_tab

    def nstar(self, image, star_groups):
        """
        Fit, as appropriate, a compound or single model to the given
        ``star_groups``. Groups are fitted sequentially from the
        smallest to the biggest. In each iteration, ``image`` is
        subtracted by the previous fitted group.

        Parameters
        ----------
        image : numpy.ndarray
            Background-subtracted image.
        star_groups : `~astropy.table.Table`
            This table must contain the following columns: ``id``,
            ``group_id``, ``x_0``, ``y_0``, ``flux_0``.  ``x_0`` and
            ``y_0`` are initial estimates of the centroids and
            ``flux_0`` is an initial estimate of the flux. Additionally,
            columns named as ``<param_name>_0`` are required if any
            other parameter in the psf model is free (i.e., the
            ``fixed`` attribute of that parameter is ``False``).

        Returns
        -------
        result_tab : `~astropy.table.Table`
            Astropy table that contains photometry results.
        image : numpy.ndarray
            Residual image.
        """

        result_tab = Table()
        for param_tab_name in self._pars_to_output.keys():
            result_tab.add_column(Column(name=param_tab_name))

        unc_tab = Table()
        for param, isfixed in self.psf_model.fixed.items():
            if not isfixed:
                unc_tab.add_column(Column(name=param + "_unc"))

        y, x = np.indices(image.shape)

        star_groups = star_groups.group_by('group_id')
        for n in range(len(star_groups.groups)):
            group_psf = get_grouped_psf_model(self.psf_model,
                                              star_groups.groups[n],
                                              self._pars_to_set)
            usepixel = np.zeros_like(image, dtype=np.bool)

            for row in star_groups.groups[n]:
                usepixel[overlap_slices(large_array_shape=image.shape,
                                        small_array_shape=self.fitshape,
                                        position=(row['y_0'], row['x_0']),
                                        mode='trim')[0]] = True

            fit_model = self.fitter(group_psf, x[usepixel], y[usepixel],
                                    image[usepixel])
            param_table = self._model_params2table(fit_model,
                                                   len(star_groups.groups[n]))
            result_tab = vstack([result_tab, param_table])

            if 'param_cov' in self.fitter.fit_info.keys():
                unc_tab = vstack([unc_tab,
                                  self._get_uncertainties(
                                      len(star_groups.groups[n]))])
            try:
                from astropy.nddata.utils import NoOverlapError
            except ImportError:
                raise ImportError("astropy 1.1 or greater is required in "
                                  "order to use this class.")
            # do not subtract if the fitting did not go well
            try:
                image = subtract_psf(image, self.psf_model, param_table,
                                     subshape=self.fitshape)
            except NoOverlapError:
                pass

        if 'param_cov' in self.fitter.fit_info.keys():
            result_tab = hstack([result_tab, unc_tab])

        return result_tab, image

    def _define_fit_param_names(self):
        """
        Convenience function to define mappings between the names of the
        columns in the initial guess table (and the name of the fitted
        parameters) and the actual name of the parameters in the model.

        This method sets the following parameters on the ``self`` object:
        * ``pars_to_set`` : Dict which maps the names of the parameters
          initial guesses to the actual name of the parameter in the
          model.
        * ``pars_to_output`` : Dict which maps the names of the fitted
          parameters to the actual name of the parameter in the model.
        """

        xname, yname, fluxname = _extract_psf_fitting_names(self.psf_model)
        self._pars_to_set = {'x_0': xname, 'y_0': yname, 'flux_0': fluxname}
        self._pars_to_output = {'x_fit': xname, 'y_fit': yname,
                                'flux_fit': fluxname}

        for p, isfixed in self.psf_model.fixed.items():
            p0 = p + '_0'
            pfit = p + '_fit'
            if p not in (xname, yname, fluxname) and not isfixed:
                self._pars_to_set[p0] = p
                self._pars_to_output[pfit] = p

    def _get_uncertainties(self, star_group_size):
        """
        Retrieve uncertainties on fitted parameters from the fitter
        object.

        Parameters
        ----------
        star_group_size : int
            Number of stars in the given group.

        Returns
        -------
        unc_tab : `~astropy.table.Table`
            Table which contains uncertainties on the fitted parameters.
            The uncertainties are reported as one standard deviation.
        """

        unc_tab = Table()
        for param_name in self.psf_model.param_names:
            if not self.psf_model.fixed[param_name]:
                unc_tab.add_column(Column(name=param_name + "_unc",
                                          data=np.empty(star_group_size)))

        if 'param_cov' in self.fitter.fit_info.keys():
            if self.fitter.fit_info['param_cov'] is not None:
                k = 0
                n_fit_params = len(unc_tab.colnames)
                for i in range(star_group_size):
                    unc_tab[i] = np.sqrt(np.diag(
                                          self.fitter.fit_info['param_cov'])
                                         )[k: k + n_fit_params]
                    k = k + n_fit_params
        return unc_tab

    def _model_params2table(self, fit_model, star_group_size):
        """
        Place fitted parameters into an astropy table.

        Parameters
        ----------
        fit_model : `astropy.modeling.Fittable2DModel` instance
            PSF or PRF model to fit the data. Could be one of the models
            in this package like `~photutils.psf.sandbox.DiscretePRF`,
            `~photutils.psf.IntegratedGaussianPRF`, or any other
            suitable 2D model.
        star_group_size : int
            Number of stars in the given group.

        Returns
        -------
        param_tab : `~astropy.table.Table`
            Table that contains the fitted parameters.
        """

        param_tab = Table()

        for param_tab_name in self._pars_to_output.keys():
            param_tab.add_column(Column(name=param_tab_name,
                                        data=np.empty(star_group_size)))

        if star_group_size > 1:
            for i in range(star_group_size):
                for param_tab_name, param_name in self._pars_to_output.items():
                    param_tab[param_tab_name][i] = getattr(fit_model,
                                                           param_name +
                                                           '_' + str(i)).value
        else:
            for param_tab_name, param_name in self._pars_to_output.items():
                param_tab[param_tab_name] = getattr(fit_model, param_name).value

        return param_tab

class IterativelySubtractedPSFPhotometry(BasicPSFPhotometry):
    """
    This class implements an iterative algorithm to perform point spread
    function photometry in crowded fields. This consists of applying a
    loop of find sources, make groups, fit groups, subtract groups, and
    then repeat until no more stars are detected or a given number of
    iterations is reached.

    Parameters
    ----------
    group_maker : callable or `~photutils.psf.GroupStarsBase`
        ``group_maker`` should be able to decide whether a given star
        overlaps with any other and label them as beloging to the same
        group.  ``group_maker`` receives as input an
        `~astropy.table.Table` object with columns named as ``id``,
        ``x_0``, ``y_0``, in which ``x_0`` and ``y_0`` have the same
        meaning of ``xcentroid`` and ``ycentroid``.  This callable must
        return an `~astropy.table.Table` with columns ``id``, ``x_0``,
        ``y_0``, and ``group_id``. The column ``group_id`` should cotain
        integers starting from ``1`` that indicate which group a given
        source belongs to. See, e.g., `~photutils.psf.DAOGroup`.
    bkg_estimator : callable, instance of any `~photutils.BackgroundBase` subclass, or None
        ``bkg_estimator`` should be able to compute either a scalar
        background or a 2D background of a given 2D image. See, e.g.,
        `~photutils.background.MedianBackground`.  If None, no
        background subtraction is performed.
    psf_model : `astropy.modeling.Fittable2DModel` instance
        PSF or PRF model to fit the data. Could be one of the models in
        this package like `~photutils.psf.sandbox.DiscretePRF`,
        `~photutils.psf.IntegratedGaussianPRF`, or any other suitable 2D
        model.  This object needs to identify three parameters (position
        of center in x and y coordinates and the flux) in order to set
        them to suitable starting values for each fit. The names of
        these parameters should be given as ``x_0``, ``y_0`` and
        ``flux``.  `~photutils.psf.prepare_psf_model` can be used to
        prepare any 2D model to match this assumption.
    fitshape : int or length-2 array-like
        Rectangular shape around the center of a star which will be used
        to collect the data to do the fitting. Can be an integer to be
        the same along both axes. E.g., 5 is the same as (5, 5), which
        means to fit only at the following relative pixel positions:
        [-2, -1, 0, 1, 2].  Each element of ``fitshape`` must be an odd
        number.
    finder : callable or instance of any `~photutils.detection.StarFinderBase` subclasses
        ``finder`` should be able to identify stars, i.e. compute a
        rough estimate of the centroids, in a given 2D image.
        ``finder`` receives as input a 2D image and returns an
        `~astropy.table.Table` object which contains columns with names:
        ``id``, ``xcentroid``, ``ycentroid``, and ``flux``. In which
        ``id`` is an integer-valued column starting from ``1``,
        ``xcentroid`` and ``ycentroid`` are center position estimates of
        the sources and ``flux`` contains flux estimates of the sources.
        See, e.g., `~photutils.detection.DAOStarFinder` or
        `~photutils.detection.IRAFStarFinder`.
    fitter : `~astropy.modeling.fitting.Fitter` instance
        Fitter object used to compute the optimized centroid positions
        and/or flux of the identified sources. See
        `~astropy.modeling.fitting` for more details on fitters.
    aperture_radius : float
        The radius (in units of pixels) used to compute initial
        estimates for the fluxes of sources. If ``None``, one FWHM will
        be used if it can be determined from the ```psf_model``.
    niters : int or None
        Number of iterations to perform of the loop FIND, GROUP,
        SUBTRACT, NSTAR. If None, iterations will proceed until no more
        stars remain.  Note that in this case it is *possible* that the
        loop will never end if the PSF has structure that causes
        subtraction to create new sources infinitely.

    Notes
    -----
    If there are problems with fitting large groups, change the
    parameters of the grouping algorithm to reduce the number of sources
    in each group or input a ``star_groups`` table that only includes
    the groups that are relevant (e.g. manually remove all entries that
    coincide with artifacts).

    References
    ----------
    [1] Stetson, Astronomical Society of the Pacific, Publications,
        (ISSN 0004-6280), vol. 99, March 1987, p. 191-222.
        Available at: http://adsabs.harvard.edu/abs/1987PASP...99..191S
    """

    def __init__(self, group_maker, bkg_estimator, psf_model, fitshape,
                 finder, fitter=LevMarLSQFitter(), niters=3,
                 aperture_radius=None):

        super().__init__(group_maker, bkg_estimator, psf_model, fitshape,
                         finder, fitter, aperture_radius)
        self.niters = niters

    @property
    def niters(self):
        return self._niters

    @niters.setter
    def niters(self, value):
        if value is None:
            self._niters = None
        else:
            try:
                if value <= 0:
                    raise ValueError('niters must be positive.')
                else:
                    self._niters = int(value)
            except ValueError:
                raise ValueError('niters must be None or an integer or '
                                 'convertable into an integer.')

    @property
    def finder(self):
        return self._finder

    @finder.setter
    def finder(self, value):
        if value is None:
            raise ValueError("finder cannot be None for "
                             "IterativelySubtractedPSFPhotometry - you may "
                             "want to use BasicPSFPhotometry. Please see the "
                             "Detection section on photutils documentation.")
        else:
            self._finder = value

    def do_photometry(self, image, init_guesses=None):
        """
        Perform PSF photometry in ``image``.

        This method assumes that ``psf_model`` has centroids and flux
        parameters which will be fitted to the data provided in
        ``image``. A compound model, in fact a sum of ``psf_model``,
        will be fitted to groups of stars automatically identified by
        ``group_maker``. Also, ``image`` is not assumed to be background
        subtracted.  If ``init_guesses`` are not ``None`` then this
        method uses ``init_guesses`` as initial guesses for the
        centroids. If the centroid positions are set as ``fixed`` in the
        PSF model ``psf_model``, then the optimizer will only consider
        the flux as a variable.

        Parameters
        ----------
        image : 2D array-like, `~astropy.io.fits.ImageHDU`, `~astropy.io.fits.HDUList`
            Image to perform photometry.
        init_guesses: `~astropy.table.Table`
            Table which contains the initial guesses (estimates) for the
            set of parameters. Columns 'x_0' and 'y_0' which represent
            the positions (in pixel coordinates) for each object must be
            present.  'flux_0' can also be provided to set initial
            fluxes.  If 'flux_0' is not provided, aperture photometry is
            used to estimate initial values for the fluxes. Additional
            columns of the form '<parametername>_0' will be used to set
            the initial guess for any parameters of the ``psf_model``
            model that are not fixed.

        Returns
        -------
        output_table : `~astropy.table.Table` or None
            Table with the photometry results, i.e., centroids and
            fluxes estimations and the initial estimates used to start
            the fitting process. Uncertainties on the fitted parameters
            are reported as columns called ``<paramname>_unc`` provided
            that the fitter object contains a dictionary called
            ``fit_info`` with the key ``param_cov``, which contains the
            covariance matrix.
        """

        if init_guesses is not None:
            table = super().do_photometry(image, init_guesses)
            table['iter_detected'] = np.ones(table['x_fit'].shape,
                                             dtype=np.int32)

            # n_start = 2 because it starts in the second iteration
            # since the first iteration is above
            output_table = self._do_photometry(init_guesses.colnames,
                                               n_start=2)
            output_table = vstack([table, output_table])
        else:
            if self.bkg_estimator is not None:
                self._residual_image = image - self.bkg_estimator(image)

            if self.aperture_radius is None:
                if hasattr(self.psf_model, 'fwhm'):
                    self.aperture_radius = self.psf_model.fwhm.value
                elif hasattr(self.psf_model, 'sigma'):
                    self.aperture_radius = (self.psf_model.sigma.value *
                                            gaussian_sigma_to_fwhm)

            output_table = self._do_photometry(['x_0', 'y_0', 'flux_0'])
        return output_table

    def _do_photometry(self, param_tab, n_start=1):
        """
        Helper function which performs the iterations of the photometry
        process.

        Parameters
        ----------
        param_names :  list
            Names of the columns which represent the initial guesses.
            For example, ['x_0', 'y_0', 'flux_0'], for intial guesses on
            the center positions and the flux.
        n_start : int
            Integer representing the start index of the iteration.  It
            is 1 if init_guesses are None, and 2 otherwise.

        Returns
        -------
        output_table : `~astropy.table.Table` or None
            Table with the photometry results, i.e., centroids and
            fluxes estimations and the initial estimates used to start
            the fitting process.
        """

        output_table = Table()
        self._define_fit_param_names()

        for (init_parname, fit_parname) in zip(self._pars_to_set.keys(),
                                               self._pars_to_output.keys()):
            output_table.add_column(Column(name=init_parname))
            output_table.add_column(Column(name=fit_parname))

        sources = self.finder(self._residual_image)

        n = n_start
        while(len(sources) > 0 and
              (self.niters is None or n <= self.niters)):
            apertures = CircularAperture((sources['xcentroid'],
                                          sources['ycentroid']),
                                         r=self.aperture_radius)
            sources['aperture_flux'] = aperture_photometry(
                self._residual_image, apertures)['aperture_sum']

            init_guess_tab = Table(names=['id', 'x_0', 'y_0', 'flux_0'],
                                   data=[sources['id'], sources['xcentroid'],
                                         sources['ycentroid'],
                                         sources['aperture_flux']])

            for param_tab_name, param_name in self._pars_to_set.items():
                if param_tab_name not in (['x_0', 'y_0', 'flux_0']):
                    init_guess_tab.add_column(
                        Column(name=param_tab_name,
                               data=(getattr(self.psf_model,
                                             param_name) *
                                     np.ones(len(sources)))))

            star_groups = self.group_maker(init_guess_tab)
            table, self._residual_image = super().nstar(
                self._residual_image, star_groups)

            star_groups = star_groups.group_by('group_id')
            table = hstack([star_groups, table])

            table['iter_detected'] = n*np.ones(table['x_fit'].shape,
                                               dtype=np.int32)

            output_table = vstack([output_table, table])

            # do not warn if no sources are found beyond the first iteration
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', AstropyUserWarning)
                sources = self.finder(self._residual_image)

            n += 1

        return output_table

class DAOPhotPSFPhotometry(IterativelySubtractedPSFPhotometry):
    """
    This class implements  an iterative algorithm based on the DAOPHOT
    algorithm presented by Stetson (1987) to perform point spread
    function photometry in crowded fields. This consists of applying a
    loop of find sources, make groups, fit groups, subtract groups, and
    then repeat until no more stars are detected or a given number of
    iterations is reached.

    Basically, this classes uses
    `~photutils.psf.IterativelySubtractedPSFPhotometry`, but with
    grouping, finding, and background estimation routines defined a
    priori. More precisely, this class uses `~photutils.psf.DAOGroup`
    for grouping, `~photutils.detection.DAOStarFinder` for finding
    sources, and `~photutils.background.MMMBackground` for background
    estimation. Those classes are based on GROUP, FIND, and SKY routines
    used in DAOPHOT, respectively.

    The parameter ``crit_separation`` is associated with
    `~photutils.psf.DAOGroup`.  ``sigma_clip`` is associated with
    `~photutils.background.MMMBackground`.  ``threshold`` and ``fwhm``
    are associated with `~photutils.detection.DAOStarFinder`.
    Parameters from ``ratio`` to ``roundhi`` are also associated with
    `~photutils.detection.DAOStarFinder`.

    Parameters
    ----------
    crit_separation : float or int
        Distance, in units of pixels, such that any two stars separated
        by less than this distance will be placed in the same group.
    threshold : float
        The absolute image value above which to select sources.
    fwhm : float
        The full-width half-maximum (FWHM) of the major axis of the
        Gaussian kernel in units of pixels.
    psf_model : `astropy.modeling.Fittable2DModel` instance
        PSF or PRF model to fit the data. Could be one of the models in
        this package like `~photutils.psf.sandbox.DiscretePRF`,
        `~photutils.psf.IntegratedGaussianPRF`, or any other suitable 2D
        model.  This object needs to identify three parameters (position
        of center in x and y coordinates and the flux) in order to set
        them to suitable starting values for each fit. The names of
        these parameters should be given as ``x_0``, ``y_0`` and
        ``flux``.  `~photutils.psf.prepare_psf_model` can be used to
        prepare any 2D model to match this assumption.
    fitshape : int or length-2 array-like
        Rectangular shape around the center of a star which will be used
        to collect the data to do the fitting. Can be an integer to be
        the same along both axes. E.g., 5 is the same as (5, 5), which
        means to fit only at the following relative pixel positions:
        [-2, -1, 0, 1, 2].  Each element of ``fitshape`` must be an odd
        number.
    sigma : float, optional
        Number of standard deviations used to perform sigma clip with a
        `astropy.stats.SigmaClip` object.
    ratio : float, optional
        The ratio of the minor to major axis standard deviations of the
        Gaussian kernel.  ``ratio`` must be strictly positive and less
        than or equal to 1.0.  The default is 1.0 (i.e., a circular
        Gaussian kernel).
    theta : float, optional
        The position angle (in degrees) of the major axis of the
        Gaussian kernel measured counter-clockwise from the positive x
        axis.
    sigma_radius : float, optional
        The truncation radius of the Gaussian kernel in units of sigma
        (standard deviation) [``1 sigma = FWHM /
        (2.0*sqrt(2.0*log(2.0)))``].
    sharplo : float, optional
        The lower bound on sharpness for object detection.
    sharphi : float, optional
        The upper bound on sharpness for object detection.
    roundlo : float, optional
        The lower bound on roundess for object detection.
    roundhi : float, optional
        The upper bound on roundess for object detection.
    fitter : `~astropy.modeling.fitting.Fitter` instance
        Fitter object used to compute the optimized centroid positions
        and/or flux of the identified sources. See
        `~astropy.modeling.fitting` for more details on fitters.
    niters : int or None
        Number of iterations to perform of the loop FIND, GROUP,
        SUBTRACT, NSTAR. If None, iterations will proceed until no more
        stars remain.  Note that in this case it is *possible* that the
        loop will never end if the PSF has structure that causes
        subtraction to create new sources infinitely.
    aperture_radius : float
        The radius (in units of pixels) used to compute initial
        estimates for the fluxes of sources. If ``None``, one FWHM will
        be used if it can be determined from the ```psf_model``.

    Notes
    -----
    If there are problems with fitting large groups, change the
    parameters of the grouping algorithm to reduce the number of sources
    in each group or input a ``star_groups`` table that only includes
    the groups that are relevant (e.g. manually remove all entries that
    coincide with artifacts).

    References
    ----------
    [1] Stetson, Astronomical Society of the Pacific, Publications,
        (ISSN 0004-6280), vol. 99, March 1987, p. 191-222.
        Available at: http://adsabs.harvard.edu/abs/1987PASP...99..191S
    """

    def __init__(self, crit_separation, threshold, fwhm, psf_model, fitshape,
                 sigma=3., ratio=1.0, theta=0.0, sigma_radius=1.5,
                 sharplo=0.2, sharphi=1.0, roundlo=-1.0, roundhi=1.0,
                 fitter=LevMarLSQFitter(),
                 niters=3, aperture_radius=None):

        self.crit_separation = crit_separation
        self.threshold = threshold
        self.fwhm = fwhm
        self.sigma = sigma
        self.ratio = ratio
        self.theta = theta
        self.sigma_radius = sigma_radius
        self.sharplo = sharplo
        self.sharphi = sharphi
        self.roundlo = roundlo
        self.roundhi = roundhi

        group_maker = DAOGroup(crit_separation=self.crit_separation)
        bkg_estimator = MMMBackground(sigma_clip=SigmaClip(sigma=self.sigma))
        finder = DAOStarFinder(threshold=self.threshold, fwhm=self.fwhm,
                               ratio=self.ratio, theta=self.theta,
                               sigma_radius=self.sigma_radius,
                               sharplo=self.sharplo, sharphi=self.sharphi,
                               roundlo=self.roundlo, roundhi=self.roundhi)

        super().__init__(group_maker=group_maker, bkg_estimator=bkg_estimator,
                         psf_model=psf_model, fitshape=fitshape,
                         finder=finder, fitter=fitter, niters=niters,
                         aperture_radius=aperture_radius)


#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import time
import ctypes
import tempfile
import numpy
import h5py
from pyscf import lib
from functools import reduce
from pyscf.lib import logger
from pyscf import gto
from pyscf import ao2mo
from pyscf.cc import ccsd
from pyscf.cc import _ccsd
from pyscf.cc import ccsd_rdm
from pyscf.scf import rhf_grad
from pyscf.scf import cphf

BLKSIZE = 192

def IX_intermediates(mycc, t1, t2, l1, l2, eris=None, d1=None, d2=None):
    if eris is None:
# Note eris are in Chemist's notation
        eris = ccsd._ERIS(mycc)
    if d1 is None:
        d1 = ccsd_rdm.gamma1_intermediates(mycc, t1, t2, l1, l2)
    doo, dov, dvo, dvv = d1
    if d2 is None:
        _d2tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
        fd2intermediate = h5py.File(_d2tmpfile.name, 'w')
        ccsd_rdm.gamma2_outcore(mycc, t1, t2, l1, l2, fd2intermediate)
        dovov = fd2intermediate['dovov']
        dvvvv = fd2intermediate['dvvvv']
        doooo = fd2intermediate['doooo']
        doovv = fd2intermediate['doovv']
        dovvo = fd2intermediate['dovvo']
        dovvv = fd2intermediate['dovvv']
        dooov = fd2intermediate['dooov']
    else:
        dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2

    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    nvir_pair = nvir * (nvir+1) //2
    _tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fswap = h5py.File(_tmpfile.name, 'w')
    fswap.create_group('e_vvov')
    fswap.create_group('c_vvov')

# Note Ioo, Ivv are not hermitian
    Ioo = numpy.zeros((nocc,nocc))
    Ivv = numpy.zeros((nvir,nvir))
    Ivo = numpy.zeros((nvir,nocc))
    Xvo = numpy.zeros((nvir,nocc))

    eris_oooo = _cp(eris.oooo)
    eris_ooov = _cp(eris.ooov)
    d_oooo = _cp(doooo)
    d_oooo = _cp(d_oooo + d_oooo.transpose(1,0,2,3))
    #:Ioo += numpy.einsum('jmlk,imlk->ij', d_oooo, eris_oooo) * 2
    Ioo += lib.dot(eris_oooo.reshape(nocc,-1), d_oooo.reshape(nocc,-1).T, 2)
    d_oooo = _cp(d_oooo.transpose(0,2,3,1))
    #:Xvo += numpy.einsum('iljk,ljka->ai', d_oooo, eris_ooov) * 2
    Xvo += lib.dot(eris_ooov.reshape(-1,nvir).T, d_oooo.reshape(nocc,-1).T, 2)
    Xvo +=(numpy.einsum('kj,kjia->ai', doo, eris_ooov) * 4
         - numpy.einsum('kj,ikja->ai', doo+doo.T, eris_ooov))
    eris_oooo = eris_ooov = d_oooo = None

    d_ovov = numpy.empty((nocc,nvir,nocc,nvir))
    blksize = 8
    for p0, p1 in prange(0, nocc, blksize):
        d_ovov[p0:p1] = _cp(dovov[p0:p1])
        d_ovvo = _cp(dovvo[p0:p1])
        for i in range(p0,p1):
            d_ovov[i] += d_ovvo[i-p0].transpose(0,2,1)
    d_ovvo = None
    d_ovov = lib.transpose_sum(d_ovov.reshape(nov,nov)).reshape(nocc,nvir,nocc,nvir)
    #:Ivo += numpy.einsum('jbka,jbki->ai', d_ovov, eris.ovoo)
    Ivo += lib.dot(d_ovov.reshape(-1,nvir).T, _cp(eris.ovoo).reshape(-1,nocc))
    eris_ovov = _cp(eris.ovov)
    #:Ioo += numpy.einsum('jakb,iakb->ij', d_ovov, eris.ovov)
    #:Ivv += numpy.einsum('jcib,jcia->ab', d_ovov, eris.ovov)
    Ioo += lib.dot(eris_ovov.reshape(nocc,-1), d_ovov.reshape(nocc,-1).T)
    Ivv += lib.dot(eris_ovov.reshape(-1,nvir).T, d_ovov.reshape(-1,nvir))
    eris_ovov = None
    fswap['dovvo'] = d_ovov.transpose(0,1,3,2)
    d_ovov = None

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = max(nvir**3*2.5, nvir**3*2+nocc*nvir**2)
    blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/unit))
    iobuflen = int(256e6/8/(blksize*nvir))
    log.debug1('IX_intermediates pass 1: block size = %d, nocc = %d in %d blocks',
               blksize, nocc, int((nocc+blksize-1)/blksize))
    for istep, (p0, p1) in enumerate(prange(0, nocc, blksize)):
        d_ooov = _cp(dooov[p0:p1])
        eris_oooo = _cp(eris.oooo[p0:p1])
        eris_ooov = _cp(eris.ooov[p0:p1])
        #:Ivv += numpy.einsum('ijkb,ijka->ab', d_ooov, eris_ooov)
        #:Ivo += numpy.einsum('jlka,jlki->ai', d_ooov, eris_oooo)
        Ivv += lib.dot(eris_ooov.reshape(-1,nvir).T, d_ooov.reshape(-1,nvir))
        Ivo += lib.dot(d_ooov.reshape(-1,nvir).T, eris_oooo.reshape(-1,nocc))
        #:Ioo += numpy.einsum('klja,klia->ij', d_ooov, eris_ooov)
        #:Xvo += numpy.einsum('kjib,kjba->ai', d_ooov, eris.oovv)
        eris_oovv = _cp(eris.oovv[p0:p1])
        tmp = _cp(d_ooov.transpose(0,1,3,2).reshape(-1,nocc))
        Ioo += lib.dot(_cp(eris_ooov.transpose(0,1,3,2).reshape(-1,nocc)).T, tmp)
        Xvo += lib.dot(eris_oovv.reshape(-1,nvir).T, tmp)
        eris_oooo = tmp = None

        d_ooov = d_ooov + dooov[:,p0:p1].transpose(1,0,2,3)
        eris_ovov = _cp(eris.ovov[p0:p1])
        #:Ioo += numpy.einsum('ljka,lika->ij', d_ooov, eris_ooov)
        #:Xvo += numpy.einsum('jikb,jakb->ai', d_ooov, eris_ovov)
        for i in range(p1-p0):
            lib.dot(eris_ooov[i].reshape(nocc,-1),
                    d_ooov[i].reshape(nocc,-1).T, 1, Ioo, 1)
            lib.dot(eris_ovov[i].reshape(nvir,-1),
                    d_ooov[i].reshape(nocc,-1).T, 1, Xvo, 1)
        d_ooov = None

        #:Ioo += numpy.einsum('kjba,kiba->ij', d_oovv, eris.oovv)
        #:Ivv += numpy.einsum('ijcb,ijca->ab', d_oovv, eris.oovv)
        #:Ivo += numpy.einsum('kjba,kjib->ai', d_oovv, eris.ooov)
        d_oovv = _cp(doovv[p0:p1]) + doovv[:,p0:p1].transpose(1,0,3,2)
        for i in range(p1-p0):
            Ioo += lib.dot(eris_oovv[i].reshape(nocc, -1), d_oovv[i].reshape(nocc,-1).T)
        Ivv += lib.dot(eris_oovv.reshape(-1,nvir).T, d_oovv.reshape(-1,nvir))
        Ivo += lib.dot(d_oovv.reshape(-1,nvir).T,
                       _cp(eris_ooov.transpose(0,1,3,2).reshape(-1,nocc)))
        eris_ooov = None
        d_oovv = _ccsd.precontract(d_oovv.reshape(-1,nvir,nvir)).reshape(p1-p0,nocc,-1)

        d_ovvv = numpy.empty((p1-p0,nvir,nvir,nvir))
        ao2mo.outcore._load_from_h5g(dovvv, p0*nvir, p1*nvir,
                                     d_ovvv.reshape(-1,nvir**2))
        #:Ivo += numpy.einsum('jadc,jidc->ai', d_ovvv, eris_oovv)
        for i in range(p1-p0):
            Ivo += lib.dot(d_ovvv[i].reshape(nvir,-1), eris_oovv[i].reshape(nocc,-1).T)
        eris_oovv = None

        # tril part of (d_ovvv + d_ovvv.transpose(0,1,3,2))
        c_ovvv = _ccsd.precontract(d_ovvv.reshape(-1,nvir,nvir))
        ao2mo.outcore._transpose_to_h5g(fswap, 'c_vvov/%d'%istep, c_ovvv, iobuflen)
        c_ovvv = c_ovvv.reshape(-1,nvir,nvir_pair)
        eris_ovx = _cp(eris.ovvv[p0:p1])
        ao2mo.outcore._transpose_to_h5g(fswap, 'e_vvov/%d'%istep,
                                        eris_ovx.reshape(-1,nvir_pair), iobuflen)
        #:Xvo += numpy.einsum('jibc,jabc->ai', d_oovv, eris_ovvv)
        #:Ivv += numpy.einsum('ibdc,iadc->ab', d_ovvv, eris_ovvv)
        for i in range(p1-p0):
            lib.dot(eris_ovx[i].reshape(nvir,-1),
                    d_oovv[i].reshape(nocc,-1).T, 1, Xvo, 1)
            lib.dot(eris_ovx[i].reshape(nvir,-1),
                    c_ovvv[i].reshape(nvir,-1).T, 1, Ivv, 1)
        c_ovvv = d_oovv = None

        eris_ovvo = numpy.empty((p1-p0,nvir,nvir,nocc))
        for i in range(p1-p0):
            d_ovvv[i] = _ccsd.sum021(d_ovvv[i])
            eris_ovvo[i] = eris_ovov[i].transpose(0,2,1)
        #:Ivo += numpy.einsum('abjc,ibjc->ai', d_ovvv, eris_ovov)
        Ivo += lib.dot(d_ovvv.reshape(-1,nvir).T, eris_ovvo.reshape(-1,nocc))
        eris_ovvo = eris_ovov = None

        eris_ovvv = lib.unpack_tril(eris_ovx.reshape(-1,nvir_pair))
        eris_ovx = None
        eris_ovvv = eris_ovvv.reshape(p1-p0,nvir,nvir,nvir)
        #:Ivv += numpy.einsum('icdb,icda->ab', d_ovvv, eris_ovvv)
        #:Xvo += numpy.einsum('jibc,jabc->ai', d_oovv, eris_ovvv)
        Ivv += lib.dot(eris_ovvv.reshape(-1,nvir).T, d_ovvv.reshape(-1,nvir))
        Xvo[:,p0:p1] +=(numpy.einsum('cb,iacb->ai', dvv, eris_ovvv) * 4
                      - numpy.einsum('cb,icba->ai', dvv+dvv.T, eris_ovvv))

        d_ovvo = _cp(fswap['dovvo'][p0:p1])
        #:Xvo += numpy.einsum('jbic,jbca->ai', d_ovov, eris_ovvv)
        lib.dot(eris_ovvv.reshape(-1,nvir).T, d_ovvo.reshape(-1,nocc), 1, Xvo, 1)

        d_ovvv = d_ovvo = eris_ovvv = None

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nocc*nvir**2 + nvir**3*2.5
    blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/unit))
    log.debug1('IX_intermediates pass 2: block size = %d, nocc = %d in %d blocks',
               blksize, nocc, int((nocc+blksize-1)/blksize))
    for p0, p1 in prange(0, nvir, blksize):
        off0 = p0*(p0+1)//2
        off1 = p1*(p1+1)//2
        d_vvvv = _cp(dvvvv[off0:off1]) * 4
        for i in range(p0, p1):
            d_vvvv[i*(i+1)//2+i-off0] *= .5
        d_vvvv = lib.unpack_tril(d_vvvv)
        eris_vvvv = lib.unpack_tril(_cp(eris.vvvv[off0:off1]))
        #:Ivv += numpy.einsum('decb,deca->ab', d_vvvv, eris_vvvv) * 2
        #:Xvo += numpy.einsum('dbic,dbca->ai', d_vvov, eris_vvvv)
        lib.dot(eris_vvvv.reshape(-1,nvir).T, d_vvvv.reshape(-1,nvir), 2, Ivv, 1)
        #:d_vvvv = _cp(d_vvvv + d_vvvv.transpose(0,1,3,2))
        d_vvov = numpy.empty((off1-off0,nocc,nvir))
        ao2mo.outcore._load_from_h5g(fswap['c_vvov'], off0, off1, d_vvov.reshape(-1,nov))
        d_vvvo = _cp(d_vvov.transpose(0,2,1))
        lib.dot(eris_vvvv.reshape(-1,nvir).T, d_vvvo.reshape(-1,nocc), 1, Xvo, 1)
        d_vvov = eris_vvvv = None

        eris_vvov = numpy.empty((off1-off0,nocc,nvir))
        ao2mo.outcore._load_from_h5g(fswap['e_vvov'], off0, off1,
                                     eris_vvov.reshape(-1,nov))
        eris_vvvo = _cp(eris_vvov.transpose(0,2,1))
        #:Ioo += numpy.einsum('abjc,abci->ij', d_vvov, eris_vvvo)
        #:Ivo += numpy.einsum('dbca,dbci->ai', d_vvvv, eris_vvvo) * 2
        lib.dot(d_vvvv.reshape(-1,nvir).T, eris_vvvo.reshape(-1,nocc), 2, Ivo, 1)
        lib.dot(eris_vvvo.reshape(-1,nocc).T, d_vvvo.reshape(-1,nocc), 1, Ioo, 1)
        eris_vvov = eris_vovv = d_vvvv = None

    del(fswap['e_vvov'])
    del(fswap['c_vvov'])
    del(fswap['dovvo'])
    fswap.close()
    _tmpfile = None

    if d2 is None:
        for key in fd2intermediate.keys():
            del(fd2intermediate[key])
        fd2intermediate.close()
        _d2tmpfile = None

    Ioo *= -1
    Ivv *= -1
    Ivo *= -1
    Xvo += Ivo
    return Ioo, Ivv, Ivo, Xvo

def response_dm1(mycc, t1, t2, l1, l2, eris=None, IX=None):
    if eris is None:
# Note eris are in Chemist's notation
        eris = ccsd._ERIS(mycc)
    if IX is None:
        Ioo, Ivv, Ivo, Xvo = IX_intermediates(mycc, t1, t2, l1, l2, eris)
    else:
        Ioo, Ivv, Ivo, Xvo = IX
    nocc, nvir = t1.shape
    nmo = nocc + nvir
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/(nocc*nvir**2)))
    def fvind(x):
        x = x.reshape(Xvo.shape)
        if eris is None:
            mo_coeff = mycc.mo_coeff
            dm = reduce(numpy.dot, (mo_coeff[:,nocc:], x, mo_coeff[:,:nocc].T))
            dm = (dm + dm.T) * 2
            v = reduce(numpy.dot, (mo_coeff[:,nocc:].T, mycc._scf.get_veff(mol, dm),
                                   mo_coeff[:,:nocc]))
        else:
            v = numpy.zeros((nocc,nvir))
            for p0, p1 in prange(0, nocc, blksize):
                eris_ovov = _cp(eris.ovov[p0:p1])
                v[p0:p1] += numpy.einsum('iajb,bj->ia', eris_ovov, x) * 4
                v[p0:p1] -= numpy.einsum('ibja,bj->ia', eris_ovov, x)
                eris_ovov = None
                v[p0:p1] -= numpy.einsum('ijba,bj->ia', _cp(eris.oovv[p0:p1]), x[:,p0:p1])
        return v.T
    mo_energy = eris.fock.diagonal()
    mo_occ = numpy.zeros_like(mo_energy)
    mo_occ[:nocc] = 2
    dvo = cphf.solve(fvind, mo_energy, mo_occ, Xvo, max_cycle=30)[0]
    dm1 = numpy.zeros((nmo,nmo))
    dm1[nocc:,:nocc] = dvo
    dm1[:nocc,nocc:] = dvo.T
    return dm1

#
# Note: only works with canonical orbitals
# Non-canonical formula refers to JCP, 95, 2639
#
def kernel(mycc, t1=None, t2=None, l1=None, l2=None, eris=None, atmlst=None,
           mf_grad=None, verbose=logger.INFO):
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if l1 is None: l1 = mycc.l1
    if l2 is None: l2 = mycc.l2
    if eris is None: eris = ccsd._ERIS(mycc)
    if mf_grad is None:
        mf_grad = rhf_grad.Gradients(mycc._scf)

    log = logger.Logger(mycc.stdout, mycc.verbose)
    time0 = time.clock(), time.time()
    mol = mycc.mol
    moidx = numpy.ones(mycc.mo_coeff.shape[1], dtype=numpy.bool)
    if isinstance(mycc.frozen, (int, numpy.integer)):
        raise NotImplementedError('frozen orbital ccsd_grad')
        moidx[:mycc.frozen] = False
    else:
        moidx[mycc.frozen] = False
    mo_coeff = mycc.mo_coeff[:,moidx]  #FIXME: ensure mycc.mo_coeff is canonical orbital
    mo_energy = eris.fock.diagonal()
    nocc, nvir = t1.shape
    nao, nmo = mo_coeff.shape
    nao_pair = nao * (nao+1) // 2

    log.debug('Build ccsd rdm1 intermediates')
    d1 = ccsd_rdm.gamma1_intermediates(mycc, t1, t2, l1, l2)
    doo, dov, dvo, dvv = d1
    time1 = log.timer('rdm1 intermediates', *time0)

    log.debug('Build ccsd rdm2 intermediates')
    _d2tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fd2intermediate = h5py.File(_d2tmpfile.name, 'w')
    d2 = ccsd_rdm.gamma2_outcore(mycc, t1, t2, l1, l2, fd2intermediate)
    time1 = log.timer('rdm2 intermediates', *time1)
    log.debug('Build ccsd response_rdm1')
    Ioo, Ivv, Ivo, Xvo = IX_intermediates(mycc, t1, t2, l1, l2, eris, d1, d2)
    time1 = log.timer('response_rdm1 intermediates', *time1)

    dm1mo = response_dm1(mycc, t1, t2, l1, l2, eris, (Ioo, Ivv, Ivo, Xvo))
    dm1mo[:nocc,:nocc] = doo + doo.T
    dm1mo[nocc:,nocc:] = dvv + dvv.T
    dm1ao = reduce(numpy.dot, (mo_coeff, dm1mo, mo_coeff.T))
    im1 = numpy.zeros_like(dm1mo)
    im1[:nocc,:nocc] = Ioo
    im1[nocc:,nocc:] = Ivv
    im1[nocc:,:nocc] = Ivo
    im1[:nocc,nocc:] = Ivo.T
    im1 = reduce(numpy.dot, (mo_coeff, im1, mo_coeff.T))
    time1 = log.timer('response_rdm1', *time1)

    log.debug('symmetrized rdm2 and MO->AO transformation')
    _dm2file = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
# Basically, 4 times of dm2 is computed. *2 in _rdm2_mo2ao, *2 in _load_block_tril
    fdm2 = h5py.File(_dm2file.name, 'w')
    dm1_with_hf = dm1mo.copy()
    for i in range(nocc):  # HF 2pdm ~ 4(ij)(kl)-2(il)(jk), diagonal+1 because of 4*dm2
        dm1_with_hf[i,i] += 1
    _rdm2_mo2ao(mycc, d2, dm1_with_hf, mo_coeff, fdm2)
    time1 = log.timer('MO->AO transformation', *time1)
    for key in fd2intermediate.keys():
        del(fd2intermediate[key])
    fd2intermediate.close()

#TODO: pass hf_grad object to compute h1 and s1
    log.debug('h1 and JK1')
    h1 = mf_grad.get_hcore(mol)
    s1 = mf_grad.get_ovlp(mol)
    zeta = lib.direct_sum('i+j->ij', mo_energy, mo_energy) * .5
    zeta[nocc:,:nocc] = mo_energy[:nocc]
    zeta[:nocc,nocc:] = mo_energy[:nocc].reshape(-1,1)
    zeta = reduce(numpy.dot, (mo_coeff, zeta*dm1mo, mo_coeff.T))
    p1 = numpy.dot(mo_coeff[:,:nocc], mo_coeff[:,:nocc].T)
    vhf4sij = reduce(numpy.dot, (p1, mycc._scf.get_veff(mol, dm1ao+dm1ao.T), p1))
    time1 = log.timer('h1 and JK1', *time1)

    # Hartree-Fock part contribution
    hf_dm1 = mycc._scf.make_rdm1(mycc._scf.mo_coeff, mycc._scf.mo_occ)
    dm1ao += hf_dm1
    zeta += mf_grad.make_rdm1e(mycc._scf.mo_energy, mycc._scf.mo_coeff,
                               mycc._scf.mo_occ)

    if atmlst is None:
        atmlst = range(mol.natm)
    offsetdic = mol.offset_nr_by_atom()
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = max(1, int(max_memory*1e6/8/(nao**3*2.5)))
    ioblksize = fdm2['dm2/0'].shape[-1]
    de = numpy.zeros((len(atmlst),3))
    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = offsetdic[ia]
# s[1] dot I, note matrix im1 is not hermitian
        de[k] =(numpy.einsum('xij,ij->x', s1[:,p0:p1], im1[p0:p1])
              + numpy.einsum('xji,ij->x', s1[:,p0:p1], im1[:,p0:p1]))
# h[1] \dot DM, *2 for +c.c.,  contribute to f1
        h1ao = mf_grad._grad_rinv(mol, ia)
        h1ao[:,p0:p1] += h1[:,p0:p1]
        de[k] +=(numpy.einsum('xij,ij->x', h1ao, dm1ao)
               + numpy.einsum('xji,ij->x', h1ao, dm1ao))
# -s[1]*e \dot DM,  contribute to f1
        de[k] -=(numpy.einsum('xij,ij->x', s1[:,p0:p1], zeta[p0:p1]  )
               + numpy.einsum('xji,ij->x', s1[:,p0:p1], zeta[:,p0:p1]))
# -vhf[s_ij[1]],  contribute to f1, *2 for s1+s1.T
        de[k] -= numpy.einsum('xij,ij->x', s1[:,p0:p1], vhf4sij[p0:p1]) * 2

# 2e AO integrals dot 2pdm
        ip0 = p0
        for b0, b1, nf in shell_prange(mol, shl0, shl1, blksize):
            eri1 = mol.intor('cint2e_ip1_sph', comp=3, aosym='s2kl',
                             shls_slice=(b0,b1,0,mol.nbas,0,mol.nbas,0,mol.nbas))
            eri1 = eri1.reshape(3,nf,nao,-1)
            dm2buf = numpy.empty((nf,nao,nao_pair))
            for ic, (i0, i1) in enumerate(prange(0, nao_pair, ioblksize)):
                _load_block_tril(fdm2['dm2/%d'%ic], ip0, ip0+nf, dm2buf[:,:,i0:i1])
            de[k] -= numpy.einsum('xijk,ijk->x', eri1, dm2buf) * 2
            eri1 = dm2buf = None
            ip0 += nf
        log.debug('grad of atom %d %s = %s', ia, mol.atom_symbol(ia), de[k])
        time1 = log.timer('grad of atom %d'%ia, *time1)

    log.note('CCSD gradinets')
    log.note('==============')
    log.note('           x                y                z')
    for k, ia in enumerate(atmlst):
        log.note('%d %s  %15.9f  %15.9f  %15.9f', ia, mol.atom_symbol(ia),
                 de[k,0], de[k,1], de[k,2])
    log.timer('CCSD gradients', *time0)
    for key in fdm2.keys():
        del(fdm2[key])
    fdm2.close()
    _d2tmpfile = _dm2file = None
    return de

def shell_prange(mol, start, stop, blksize):
    nao = 0
    ib0 = start
    for ib in range(start, stop):
        now = (mol.bas_angular(ib)*2+1) * mol.bas_nctr(ib)
        nao += now
        if nao > blksize and nao > now:
            yield (ib0, ib, nao-now)
            ib0 = ib
            nao = now
    yield (ib0, stop, nao)

def _rdm2_mo2ao(mycc, d2, dm1, mo_coeff, fsave=None):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    if fsave is None:
        _dm2file = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
        fsave = h5py.File(_dm2file.name, 'w')
    else:
        _dm2file = None
    time1 = time.clock(), time.time()
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
    nocc, nvir = dovov.shape[:2]
    nov = nocc * nvir
    nao, nmo = mo_coeff.shape
    nao_pair = nao * (nao+1) // 2
    nvir_pair = nvir * (nvir+1) //2
    mo_coeff = numpy.asarray(mo_coeff, order='F')
    def _trans(vin, orbs_slice, out=None):
        nrow = vin.shape[0]
        if out is None:
            out = numpy.empty((nrow,nao_pair))
        fdrv = getattr(_ccsd.libcc, 'AO2MOnr_e2_drv')
        pao_loc = ctypes.POINTER(ctypes.c_void_p)()
        fdrv(_ccsd.libcc.AO2MOtranse2_nr_s1, _ccsd.libcc.CCmmm_transpose_sum,
             out.ctypes.data_as(ctypes.c_void_p),
             vin.ctypes.data_as(ctypes.c_void_p),
             mo_coeff.ctypes.data_as(ctypes.c_void_p),
             ctypes.c_int(nrow), ctypes.c_int(nao),
             (ctypes.c_int*4)(*orbs_slice), pao_loc, ctypes.c_int(0))
        return out

# transform dm2_ij to get lower triangular (dm2+dm2.transpose(0,1,3,2))
    _tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fswap = h5py.File(_tmpfile.name)
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = max(1, int(max_memory*1e6/8/(nmo*nao_pair+nmo**3+nvir**3)))
    iobuflen = int(256e6/8/(blksize*nmo))
    log.debug1('_rdm2_mo2ao pass 1: blksize = %d, iobuflen = %d', blksize, iobuflen)
    fswap.create_group('o')  # for h5py old version
    pool1 = numpy.empty((blksize,nmo,nmo,nmo))
    pool2 = numpy.empty((blksize,nmo,nao_pair))
    bufd_ovvv = numpy.empty((blksize,nvir,nvir,nvir))
    for istep, (p0, p1) in enumerate(prange(0, nocc, blksize)):
        buf1 = pool1[:p1-p0]
        buf1[:,:nocc,:nocc,:nocc] = doooo[p0:p1]
        buf1[:,:nocc,:nocc,nocc:] = dooov[p0:p1]
        buf1[:,:nocc,nocc:,:nocc] = 0
        buf1[:,:nocc,nocc:,nocc:] = doovv[p0:p1]
        buf1[:,nocc:,:nocc,:nocc] = 0
        buf1[:,nocc:,:nocc,nocc:] = dovov[p0:p1]
        buf1[:,nocc:,nocc:,:nocc] = dovvo[p0:p1]
        d_ovvv = bufd_ovvv[:p1-p0]
        ao2mo.outcore._load_from_h5g(dovvv, p0*nvir, p1*nvir,
                                     d_ovvv.reshape(-1,nvir**2))
        buf1[:,nocc:,nocc:,nocc:] = d_ovvv
        for i in range(p0, p1):
            buf1[i-p0,i,:,:] += dm1
            buf1[i-p0,:,:,i] -= dm1 * .5
        buf2 = pool2[:p1-p0].reshape(-1,nao_pair)
        _trans(buf1.reshape(-1,nmo**2), (0,nmo,0,nmo), buf2)
        ao2mo.outcore._transpose_to_h5g(fswap, 'o/%d'%istep, buf2, iobuflen)
    pool1 = pool2 = bufd_ovvv = None
    time1 = log.timer_debug1('_rdm2_mo2ao pass 1', *time1)

    fswap.create_group('v')  # for h5py old version
    pool1 = numpy.empty((blksize*nvir,nao_pair))
    pool2 = numpy.empty((blksize*nvir,nvir,nvir))
    for istep, (p0, p1) in enumerate(prange(0, nvir_pair, blksize*nvir)):
        buf1 = _cp(dvvvv[p0:p1])
        buf2 = lib.unpack_tril(buf1, out=pool2[:p1-p0])
        buf1 = _trans(buf2, (nocc,nmo,nocc,nmo), out=pool1[:p1-p0])
        ao2mo.outcore._transpose_to_h5g(fswap, 'v/%d'%istep, buf1, iobuflen)
    pool1 = pool2 = None
    time1 = log.timer_debug1('_rdm2_mo2ao pass 2', *time1)

# transform dm2_kl then dm2 + dm2.transpose(2,3,0,1)
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = max(nao, int(max_memory*1e6/8/(nao_pair+nmo**2)))
    iobuflen = int(256e6/8/blksize)
    log.debug1('_rdm2_mo2ao pass 3: blksize = %d, iobuflen = %d', blksize, iobuflen)
    gsave = fsave.create_group('dm2')
    for istep, (p0, p1) in enumerate(prange(0, nao_pair, blksize)):
        gsave.create_dataset(str(istep), (nao_pair,p1-p0), 'f8')
    diagidx = numpy.arange(nao)
    diagidx = diagidx*(diagidx+1)//2 + diagidx
    pool1 = numpy.empty((blksize,nmo,nmo))
    pool2 = numpy.empty((blksize,nvir_pair))
    pool3 = numpy.empty((blksize,nvir,nvir))
    pool4 = numpy.empty((blksize,nao_pair))
    for istep, (p0, p1) in enumerate(prange(0, nao_pair, blksize)):
        buf1 = pool1[:p1-p0]
        ao2mo.outcore._load_from_h5g(fswap['o'], p0, p1,
                                     buf1[:,:nocc].reshape(p1-p0,-1))
        buf2 = ao2mo.outcore._load_from_h5g(fswap['v'], p0, p1, pool2[:p1-p0])
        buf3 = lib.unpack_tril(buf2, out=pool3[:p1-p0])
        buf1[:,nocc:,nocc:] = buf3
        buf1[:,nocc:,:nocc] = 0
        buf2 = _trans(buf1, (0,nmo,0,nmo), out=pool4[:p1-p0])
        ic = 0
        idx = diagidx[diagidx<p1]
        if p0 > 0:
            buf1 = _cp(gsave[str(istep)][:p0])
            for i0, i1 in prange(0, p1-p0, BLKSIZE):
                for j0, j1, in prange(0, p0, BLKSIZE):
                    buf1[j0:j1,i0:i1] += buf2[i0:i1,j0:j1].T
                    buf2[i0:i1,j0:j1] = buf1[j0:j1,i0:i1].T
            buf1[:,idx[p0<=idx]-p0] *= .5
            gsave[str(istep)][:p0] = buf1
        lib.transpose_sum(buf2[:,p0:p1], inplace=True)
        buf2[:,idx] *= .5
        for ic, (i0, i1) in enumerate(prange(0, nao_pair, blksize)):
            gsave[str(ic)][p0:p1] = buf2[:,i0:i1]
    time1 = log.timer_debug1('_rdm2_mo2ao pass 3', *time1)
    del(fswap['o'])
    del(fswap['v'])
    fswap.close()
    _tmpfile = None
    time1 = log.timer_debug1('_rdm2_mo2ao cleanup', *time1)
    if _dm2file is not None:
        nvir_pair = nvir * (nvir+1) // 2
        dm2 = numpy.empty((nvir_pair, nvir_pair))
        ao2mo.outcore._load_from_h5g(fsave['dm2'], 0, nvir_pair, dm2)
        fsave.close()
        _dm2file = None
        return dm2
    else:
        return fsave

#
# .
# . .
# ----+             -----------
# ----|-+       =>  -----------
# . . | | .
# . . | | . .
#
def _load_block_tril(dat, row0, row1, out=None):
    shape = dat.shape
    nd = int(numpy.sqrt(shape[0]*2))
    if out is None:
        out = numpy.empty((row1-row0,nd)+shape[1:])
    else:
        out = numpy.ndarray((row1-row0,nd)+shape[1:], buffer=out)
    p0 = row0*(row0+1)//2
    for i in range(row0, row1):
        out[i-row0,:i+1] = _cp(dat[p0:p0+i+1])
        for j in range(row0, i):
            out[j-row0,i] = out[i-row0,j]
        p0 += i + 1
    for i in range(row1, nd):
        i2 = i*(i+1)//2
        out[:,i] = dat[i2+row0:i2+row1]
    return out

def hf_get_jk_incore(eri, dm):
    ni, nj = eri.shape[:2]
    vj = numpy.empty((ni,nj))
    vk = numpy.empty((ni,nj))
    _ccsd.libcc.CCvhfs2kl(eri.ctypes.data_as(ctypes.c_void_p),
                          dm.ctypes.data_as(ctypes.c_void_p),
                          vj.ctypes.data_as(ctypes.c_void_p),
                          vk.ctypes.data_as(ctypes.c_void_p),
                          ctypes.c_int(ni), ctypes.c_int(nj))
    return vj, vk

def prange(start, end, step):
    for i in range(start, end, step):
        yield i, min(i+step, end)

def _cp(a):
    return numpy.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo
    from pyscf import grad

    mol = gto.M()
    mf = scf.RHF(mol)

    mycc = ccsd.CCSD(mf)

    numpy.random.seed(2)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*20
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)

    h1 = fock0 - (numpy.einsum('kkpq->pq', eri0[:nocc,:nocc])*2
                - numpy.einsum('pkkq->pq', eri0[:,:nocc,:nocc]))
    eris = lambda:None
    idx = numpy.tril_indices(nvir)
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovo = eri0[:nocc,:nocc,nocc:,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:]
    eris.ovvv = eris.ovvv[:,:,idx[0],idx[1]].copy()
    eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:]
    eris.vvvv = eris.vvvv[idx[0],idx[1]][:,idx[0],idx[1]].copy()
    eris.fock = fock0

    print('-----------------------------------')
    Ioo, Ivv, Ivo, Xvo = IX_intermediates(mycc, t1, t2, l1, l2, eris)
    numpy.random.seed(1)
    h1 = numpy.random.random((nmo,nmo))
    h1 = h1 + h1.T
    print(numpy.einsum('ij,ij', h1[:nocc,:nocc], Ioo) - 2613213.0346526774)
    print(numpy.einsum('ab,ab', h1[nocc:,nocc:], Ivv) - 6873038.9907923322)
    print(numpy.einsum('ai,ai', h1[nocc:,:nocc], Ivo) - 4353360.4241635408)
    print(numpy.einsum('ai,ai', h1[nocc:,:nocc], Xvo) - 203575.42337558540)
    dm1 = response_dm1(mycc, t1, t2, l1, l2, eris)
    print(numpy.einsum('pq,pq', h1[nocc:,:nocc], dm1[nocc:,:nocc])--486.638981725713393)

    print('-----------------------------------')
    mol = gto.M(
        verbose = 0,
        atom = [
            ["O" , (0. , 0.     , 0.)],
            [1   , (0. ,-0.757  , 0.587)],
            [1   , (0. , 0.757  , 0.587)]],
        basis = '631g'
    )
    mf = scf.RHF(mol)
    ehf = mf.scf()

    mycc = ccsd.CCSD(mf)
    mycc.conv_tol = 1e-10
    mycc.conv_tol_normt = 1e-10
    ecc, t1, t2 = mycc.kernel()
    l1, l2 = mycc.solve_lambda()
    g1 = kernel(mycc, t1, t2, l1, l2, mf_grad=grad.RHF(mf))
    print('gcc')
    print(g1 + grad.grad_nuc(mol))
#[[ 0   0                1.00950925e-02]
# [ 0   2.28063426e-02  -5.04754623e-03]
# [ 0  -2.28063426e-02  -5.04754623e-03]]

    lib.parameters.BOHR = 1
    r = 1.76#.748
    mol = gto.M(
        verbose = 0,
        atom = '''H 0 0 0; H 0 0 %f''' % r,
        basis = '631g')
    mf = scf.RHF(mol)
    mf.conv_tol = 1e-14
    ehf0 = mf.scf()
    ghf = grad.RHF(mf).grad()
    mycc = ccsd.CCSD(mf)
    mycc.conv_tol = 1e-10
    mycc.conv_tol_normt = 1e-10
    ecc, t1, t2 = mycc.kernel()
    l1, l2 = mycc.solve_lambda()
    g1 = kernel(mycc, t1, t2, l1, l2, mf_grad=grad.RHF(mf))
    print('gcc')
    print(g1 + grad.grad_nuc(mol))
#[[ 0.          0.         -0.07080036]
# [ 0.          0.          0.07080036]]