# Adversarial learning for event-based music generation with SeqGAN
# Reference:
# "SeqGAN: Sequence Generative Adversarial Nets with Policy Gradient."
# (Yu, Lantao, et al.).
# ... Honestly, it's too hard to train ;(

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical

import numpy as np
import os, sys, time, argparse
from progress.bar import Bar

import config, utils
from config import device
from data import Dataset
from model import PerformanceRNN
from sequence import EventSeq, ControlSeq

# pylint: disable=E1101

#========================================================================
# Discriminator
#========================================================================

discriminator_config = {
    'event_dim': EventSeq.dim(),
    'hidden_dim': 512,
    'gru_layers': 3,
    'gru_dropout': 0.3
}

class EventSequenceEncoder(nn.Module):
    def __init__(self, event_dim=EventSeq.dim(), hidden_dim=512,
                 gru_layers=3, gru_dropout=0.3):
        super().__init__()
        self.event_embedding = nn.Embedding(event_dim, hidden_dim)
        self.gru = nn.GRU(hidden_dim, hidden_dim,
                          num_layers=gru_layers, dropout=gru_dropout)
        self.attn = nn.Parameter(torch.randn(hidden_dim), requires_grad=True)
        self.output_fc = nn.Linear(hidden_dim, 1)
        self.output_fc_activation = nn.Sigmoid()

    def forward(self, events, hidden=None, output_logits=False):
        # events: [steps, batch_size]
        events = self.event_embedding(events)
        outputs, _ = self.gru(events, hidden) # [t, b, h]
        weights = (outputs * self.attn).sum(-1, keepdim=True)
        output = (outputs * weights).mean(0) # [b, h]
        output = self.output_fc(output).squeeze(-1) # [b]
        if output_logits:
            return output
        output = self.output_fc_activation(output)
        return output

#========================================================================
# Pretrain Discriminator
#========================================================================

def pretrain_discriminator(model_sess_path,         # load
                           discriminator_sess_path, # load + save
                           batch_data_generator,    # Dataset(...).batches(...)
                           discriminator_config_overwrite={},
                           gradient_clipping=False,
                           control_ratio=1.0,
                           num_iter=-1,
                           save_interval=60.0,
                           discriminator_lr=0.001,
                           enable_logging=False,
                           auto_sample_factor=False,
                           sample_factor=1.0):

    print('-' * 70)
    print('model_sess_path:', model_sess_path)
    print('discriminator_sess_path:', discriminator_sess_path)
    print('discriminator_config_overwrite:', discriminator_config_overwrite)
    print('sample_factor:', sample_factor)
    print('auto_sample_factor:', auto_sample_factor)
    print('discriminator_lr:', discriminator_lr)
    print('gradient_clipping:', gradient_clipping)
    print('control_ratio:', control_ratio)
    print('num_iter:', num_iter)
    print('save_interval:', save_interval)
    print('enable_logging:', enable_logging)
    print('-' * 70)

    # Load generator
    model_sess = torch.load(model_sess_path)
    model_config = model_sess['model_config']
    model = PerformanceRNN(**model_config).to(device)
    model.load_state_dict(model_sess['model_state'])

    print(f'Generator from "{model_sess_path}"')
    print(model)
    print('-' * 70)

    # Load discriminator and optimizer
    global discriminator_config
    try:
        discriminator_sess = torch.load(discriminator_sess_path)
        discriminator_config = discriminator_sess['discriminator_config']
        discriminator_state = discriminator_sess['discriminator_state']
        discriminator_optimizer_state = discriminator_sess['discriminator_optimizer_state']
        print(f'Discriminator from "{discriminator_sess_path}"')
        discriminator_loaded = True
    except:
        print(f'New discriminator session at "{discriminator_sess_path}"')
        discriminator_config.update(discriminator_config_overwrite)
        discriminator_loaded = False

    discriminator = EventSequenceEncoder(**discriminator_config).to(device)
    optimizer = optim.Adam(discriminator.parameters(), lr=discriminator_lr)
    if discriminator_loaded:
        discriminator.load_state_dict(discriminator_state)
        optimizer.load_state_dict(discriminator_optimizer_state)

    print(discriminator)
    print(optimizer)
    print('-' * 70)

    def save_discriminator():
        print(f'Saving to "{discriminator_sess_path}"')
        torch.save({
            'discriminator_config': discriminator_config,
            'discriminator_state': discriminator.state_dict(),
            'discriminator_optimizer_state': optimizer.state_dict()
        }, discriminator_sess_path)
        print('Done saving')

    # Disable gradient for generator
    for parameter in model.parameters():
        parameter.requires_grad_(False)

    model.eval()
    discriminator.train()

    loss_func = nn.BCEWithLogitsLoss()
    last_save_time = time.time()

    if enable_logging:
        from tensorboardX import SummaryWriter
        writer = SummaryWriter()

    try:
        for i, (events, controls) in enumerate(batch_data_generator):
            if i == num_iter:
                break

            steps, batch_size = events.shape

            # Prepare inputs
            events = torch.LongTensor(events).to(device)
            if np.random.random() <= control_ratio:
                controls = torch.FloatTensor(controls).to(device)
            else:
                controls = None

            init = torch.randn(batch_size, model.init_dim).to(device)

            # Predict for real event sequence
            real_events = events
            real_logit = discriminator(real_events, output_logits=True)
            real_target = torch.ones_like(real_logit).to(device)

            if auto_sample_factor:
                sample_factor = np.random.choice([
                    0.1, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0,
                    1.1, 1.2, 1.4, 1.6, 2.0, 4.0, 10.0])

            # Predict for fake event sequence from the generator
            fake_events = model.generate(init, steps, None, controls,
                                         greedy=0, output_type='index',
                                         temperature=sample_factor)
            fake_logit = discriminator(fake_events, output_logits=True)
            fake_target = torch.zeros_like(fake_logit).to(device)

            # Compute loss
            loss = (loss_func(real_logit, real_target) +
                    loss_func(fake_logit, fake_target)) / 2

            # Backprop
            discriminator.zero_grad()
            loss.backward()

            # Gradient clipping
            norm = utils.compute_gradient_norm(discriminator.parameters())
            if gradient_clipping:
                nn.utils.clip_grad_norm_(discriminator.parameters(), gradient_clipping)

            optimizer.step()

            # Logging
            loss = loss.item()
            norm = norm.item()
            print(f'{i} loss: {loss}, norm: {norm}, sf: {sample_factor}')
            if enable_logging:
                writer.add_scalar(f'pretrain/D/loss/all', loss, i)
                writer.add_scalar(f'pretrain/D/loss/{sample_factor}', loss, i)
                writer.add_scalar(f'pretrain/D/norm/{sample_factor}', norm, i)

            if last_save_time + save_interval < time.time():
                last_save_time = time.time()
                save_discriminator()

    except KeyboardInterrupt:
        save_discriminator()

#========================================================================
# Adversarial Learning
#========================================================================

def train_adversarial(sess_path, batch_data_generator,
                      model_load_path, model_optimizer_class,
                      model_gradient_clipping, discriminator_gradient_clipping,
                      model_learning_rate, reset_model_optimizer,
                      discriminator_load_path, discriminator_optimizer_class,
                      discriminator_learning_rate, reset_discriminator_optimizer,
                      g_max_q_mean, g_min_q_mean, d_min_loss, g_max_steps, d_max_steps,
                      mc_sample_size, mc_sample_factor, first_to_train,
                      save_interval, control_ratio, enable_logging):

    if enable_logging:
        from tensorboardX import SummaryWriter
        writer = SummaryWriter()

    if os.path.isfile(sess_path):
        adv_state = torch.load(sess_path)
        model_config = adv_state['model_config']
        model_state = adv_state['model_state']
        model_optimizer_state = adv_state['model_optimizer_state']
        discriminator_config = adv_state['discriminator_config']
        discriminator_state = adv_state['discriminator_state']
        discriminator_optimizer_state = adv_state['discriminator_optimizer_state']
        print('-' * 70)
        print('Session is loaded from', sess_path)
        loaded_from_session = True

    else:
        model_sess = torch.load(model_load_path)
        model_config = model_sess['model_config']
        model_state = model_sess['model_state']
        discriminator_sess = torch.load(discriminator_load_path)
        discriminator_config = discriminator_sess['discriminator_config']
        discriminator_state = discriminator_sess['discriminator_state']
        loaded_from_session = False

    model = PerformanceRNN(**model_config)
    model.load_state_dict(model_state)
    model.to(device).train()
    model_optimizer = model_optimizer_class(model.parameters(), lr=model_learning_rate)

    discriminator = EventSequenceEncoder(**discriminator_config)
    discriminator.load_state_dict(discriminator_state)
    discriminator.to(device).train()
    discriminator_optimizer = discriminator_optimizer_class(discriminator.parameters(),
                                                            lr=discriminator_learning_rate)

    if loaded_from_session:
        if not reset_model_optimizer:
            model_optimizer.load_state_dict(model_optimizer_state)
        if not reset_discriminator_optimizer:
            discriminator_optimizer.load_state_dict(discriminator_optimizer_state)

    g_loss_func = nn.CrossEntropyLoss()
    d_loss_func = nn.BCEWithLogitsLoss(reduce=False)

    print('-' * 70)
    print('Options')
    print('sess_path:', sess_path)
    print('save_interval:', save_interval)
    print('batch_data_generator:', batch_data_generator)
    print('control_ratio:', control_ratio)
    print('g_max_q_mean:', g_max_q_mean)
    print('g_min_q_mean:', g_min_q_mean)
    print('d_min_loss:', d_min_loss)
    print('mc_sample_size:', mc_sample_size)
    print('mc_sample_factor:', mc_sample_factor)
    print('enable_logging:', enable_logging)
    print('model_load_path:', model_load_path)
    print('model_loss:', g_loss_func)
    print('model_optimizer_class:', model_optimizer_class)
    print('model_gradient_clipping:', model_gradient_clipping)
    print('model_learning_rate:', model_learning_rate)
    print('reset_model_optimizer:', reset_model_optimizer)
    print('discriminator_load_path:', discriminator_load_path)
    print('discriminator_loss:', d_loss_func)
    print('discriminator_optimizer_class:', discriminator_optimizer_class)
    print('discriminator_gradient_clipping:', discriminator_gradient_clipping)
    print('discriminator_learning_rate:', discriminator_learning_rate)
    print('reset_discriminator_optimizer:', reset_discriminator_optimizer)
    print('first_to_train:', first_to_train)
    print('-' * 70)
    print(f'Generator from "{sess_path if loaded_from_session else model_load_path}"')
    print(model)
    print(model_optimizer)
    print('-' * 70)
    print(f'Discriminator from "{sess_path if loaded_from_session else discriminator_load_path}"')
    print(discriminator)
    print(discriminator_optimizer)
    print('-' * 70)

    def save():
        print(f'Saving to "{sess_path}"')
        torch.save({
            'model_config': model_config,
            'model_state': model.state_dict(),
            'model_optimizer_state': model_optimizer.state_dict(),
            'discriminator_config': discriminator_config,
            'discriminator_state': discriminator.state_dict(),
            'discriminator_optimizer_state': discriminator_optimizer.state_dict()
        }, sess_path)
        print('Done saving')

    def mc_rollout(generated, hidden, total_steps, controls=None):
        # generated: [t, batch_size]
        # hidden: [n_layers, batch_size, hidden_dim]
        # controls: [total_steps - t, batch_size, control_dim]
        generated = torch.cat(generated, 0)
        generated_steps, batch_size = generated.shape # t, b
        steps = total_steps - generated_steps # s

        generated = generated.unsqueeze(1) # [t, 1, b]
        generated = generated.repeat(1, mc_sample_size, 1) # [t, mcs, b]
        generated = generated.view(generated_steps, -1) # [t, mcs * b]

        hidden = hidden.unsqueeze(1).repeat(1, mc_sample_size, 1, 1)
        hidden = hidden.view(model.gru_layers, -1, model.hidden_dim)

        if controls is not None:
            assert controls.shape == (steps, batch_size, model.control_dim)
            controls = controls.unsqueeze(1) # [s, 1, b, c]
            controls = controls.repeat(1, mc_sample_size, 1, 1) # [s, mcs, b, c]
            controls = controls.view(steps, -1, model.control_dim) # [s, mcs * b, c]

        event = generated[-1].unsqueeze(0) # [1, mcs * b]
        control = None # default when controls is None
        outputs = []

        for i in range(steps):
            if controls is not None:
                control = controls[i].unsqueeze(0) # [1, mcs * b, c]

            output, hidden = model.forward(event, control=control, hidden=hidden)
            probs = model.output_fc_activation(output / mc_sample_factor)
            event = Categorical(probs).sample() # [1, mcs * b]
            outputs.append(event)

        sequences = torch.cat([generated, *outputs], 0)
        assert sequences.shape == (total_steps, mc_sample_size * batch_size)
        return sequences

    def train_generator(batch_size, init, events, controls):
        # Generator step
        hidden = model.init_to_hidden(init)
        event = model.get_primary_event(batch_size)
        outputs = []
        generated = []
        q_values = []

        for step in Bar('MC Rollout').iter(range(steps)):
            control = controls[step].unsqueeze(0) if use_control else None
            output, hidden = model.forward(event, control=control, hidden=hidden)
            outputs.append(output)
            probs = model.output_fc_activation(output / mc_sample_factor)
            generated.append(Categorical(probs).sample())

            with torch.no_grad():
                if step < steps - 1:
                    sequences = mc_rollout(generated, hidden, steps, controls[step+1:])
                    mc_score = discriminator(sequences) # [mcs * b]
                    mc_score = mc_score.view(mc_sample_size, batch_size) # [mcs, b]
                    q_value = mc_score.mean(0, keepdim=True) # [1, batch_size]

                else:
                    q_value = discriminator(torch.cat(generated, 0))
                    q_value = q_value.unsqueeze(0) # [1, batch_size]

                q_values.append(q_value)

        # Compute loss
        q_values = torch.cat(q_values, 0) # [steps, batch_size]
        q_mean = q_values.mean().detach()
        q_values = q_values - q_mean
        generated = torch.cat(generated, 0) # [steps, batch_size]
        outputs = torch.cat(outputs, 0) # [steps, batch_size, event_dim]
        loss = F.cross_entropy(outputs.view(-1, model.event_dim),
                               generated.view(-1),
                               reduce=False)
        loss = (loss * q_values.view(-1)).mean()

        # Backprop
        model.zero_grad()
        loss.backward()

        # Gradient clipping
        norm = utils.compute_gradient_norm(model.parameters())
        if model_gradient_clipping:
            nn.utils.clip_grad_norm_(model.parameters(), model_gradient_clipping)

        model_optimizer.step()

        q_mean = q_mean.item()
        norm = norm.item()
        return q_mean, norm

    def train_discriminator(batch_size, init, events, controls):
        # Discriminator step
        with torch.no_grad():
            generated = model.generate(init, steps, None, controls,
                                        greedy=0, temperature=mc_sample_factor)

        fake_logit = discriminator(generated, output_logits=True)
        real_logit = discriminator(events, output_logits=True)
        fake_target = torch.zeros_like(fake_logit)
        real_target = torch.ones_like(real_logit)

        # Compute loss
        fake_loss = F.binary_cross_entropy_with_logits(fake_logit, fake_target)
        real_loss = F.binary_cross_entropy_with_logits(real_logit, real_target)
        loss = (real_loss + fake_loss) / 2

        # Backprop
        discriminator.zero_grad()
        loss.backward()

        # Gradient clipping
        norm = utils.compute_gradient_norm(discriminator.parameters())
        if discriminator_gradient_clipping:
            nn.utils.clip_grad_norm_(discriminator.parameters(), discriminator_gradient_clipping)

        discriminator_optimizer.step()

        real_loss = real_loss.item()
        fake_loss = fake_loss.item()
        loss = loss.item()
        norm = norm.item()
        return loss, real_loss, fake_loss, norm

    try:
        last_save_time = time.time()
        step_for = first_to_train
        g_steps = 0
        d_steps = 0

        for i, (events, controls) in enumerate(batch_data_generator):
            steps, batch_size = events.shape
            init = torch.randn(batch_size, model.init_dim).to(device)
            events = torch.LongTensor(events).to(device)

            use_control = np.random.random() <= control_ratio
            controls = torch.FloatTensor(controls).to(device) if use_control else None

            if step_for == 'G':
                q_mean, norm = train_generator(batch_size, init, events, controls)
                g_steps += 1

                print(f'{i} (G-step) Q_mean: {q_mean}, norm: {norm}')
                if enable_logging:
                    writer.add_scalar('adversarial/G/Q_mean', q_mean, i)
                    writer.add_scalar('adversarial/G/norm', norm, i)

                if q_mean < g_min_q_mean:
                    print(f'Q is too small: {q_mean}, exiting')
                    raise KeyboardInterrupt

                if q_mean > g_max_q_mean or (g_max_steps and g_steps >= g_max_steps):
                    step_for = 'D'
                    d_steps = 0

            if step_for == 'D':
                loss, real_loss, fake_loss, norm = train_discriminator(batch_size, init, events, controls)
                d_steps += 1

                print(f'{i} (D-step) loss: {loss} (real: {real_loss}, fake: {fake_loss}), norm: {norm}')
                if enable_logging:
                    writer.add_scalar('adversarial/D/loss', loss, i)
                    writer.add_scalar('adversarial/D/norm', norm, i)

                if fake_loss <= real_loss < d_min_loss or (d_max_steps and d_steps >= d_max_steps):
                    step_for = 'G'
                    g_steps = 0

            if last_save_time + save_interval < time.time():
                last_save_time = time.time()
                save()

    except KeyboardInterrupt:
        save()

#========================================================================
# Script Arguments
#========================================================================

def batch_generator(args):
    print('-' * 70)
    dataset = Dataset(args.dataset_path, verbose=True)
    print(dataset)
    return dataset.batches(args.batch_size, args.window_size, args.stride_size)

def pretrain(args):
    pretrain_discriminator(model_sess_path=args.generator_session_path,
                           discriminator_sess_path=args.discriminator_session_path,
                           discriminator_config_overwrite=utils.params2dict(args.discriminator_parameters),
                           batch_data_generator=args.batch_generator(args),
                           gradient_clipping=args.gradient_clipping,
                           sample_factor=args.sample_factor,
                           auto_sample_factor=args.auto_sample_factor,
                           control_ratio=args.control_ratio,
                           num_iter=args.stop_iteration,
                           save_interval=args.save_interval,
                           discriminator_lr=args.discriminator_learning_rate,
                           enable_logging=args.enable_logging)

def adversarial(args):
    train_adversarial(sess_path=args.session_path,
                      batch_data_generator=args.batch_generator(args),
                      model_load_path=args.generator_load_path,
                      discriminator_load_path=args.discriminator_load_path,
                      model_optimizer_class=getattr(optim, args.generator_optimizer),
                      discriminator_optimizer_class=getattr(optim, args.discriminator_optimizer),
                      model_gradient_clipping=args.generator_gradient_clipping,
                      discriminator_gradient_clipping=args.discriminator_gradient_clipping,
                      model_learning_rate=args.generator_learning_rate,
                      discriminator_learning_rate=args.discriminator_learning_rate,
                      reset_model_optimizer=args.reset_generator_optimizer,
                      reset_discriminator_optimizer=args.reset_discriminator_optimizer,
                      g_max_q_mean=args.g_max_q_mean,
                      g_min_q_mean=args.g_min_q_mean,
                      d_min_loss=args.d_min_loss,
                      g_max_steps=args.g_max_steps,
                      d_max_steps=args.d_max_steps,
                      mc_sample_size=args.monte_carlo_sample_size,
                      mc_sample_factor=args.monte_carlo_sample_factor,
                      first_to_train=args.first_to_train,
                      control_ratio=args.control_ratio,
                      save_interval=args.save_interval,
                      enable_logging=args.enable_logging)

def get_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()
    parser.add_argument('-d', '--dataset-path', type=str, required=True)
    parser.add_argument('-b', '--batch-size', type=int, default=64)
    parser.add_argument('-w', '--window-size', type=int, default=200)
    parser.add_argument('-s', '--stride-size', type=int, default=10)
    parser.set_defaults(batch_generator=batch_generator)
    pre_parser = subparsers.add_parser('pretrain', aliases=['p', 'pre'])
    pre_parser.add_argument('-G', '--generator-session-path', type=str, default=True)
    pre_parser.add_argument('-D', '--discriminator-session-path', type=str, required=True)
    pre_parser.add_argument('-p', '--discriminator-parameters', type=str, default='')
    pre_parser.add_argument('-l', '--discriminator-learning-rate', type=float, default=0.001)
    pre_parser.add_argument('-g', '--gradient-clipping', type=float, default=False)
    pre_parser.add_argument('-f', '--sample-factor', type=float, default=1.0)
    pre_parser.add_argument('-af', '--auto-sample-factor', action='store_true', default=False)
    pre_parser.add_argument('-c', '--control-ratio', type=float, default=1.0)
    pre_parser.add_argument('-n', '--stop-iteration', type=int, default=-1)
    pre_parser.add_argument('-i', '--save-interval', type=float, default=60.0)
    pre_parser.add_argument('-L', '--enable-logging', action='store_true', default=False)
    pre_parser.set_defaults(main=pretrain)
    adv_parser = subparsers.add_parser('adversarial', aliases=['a', 'adv'])
    adv_parser.add_argument('-S', '--session-path', type=str, required=True)
    adv_parser.add_argument('-Gp', '--generator-load-path', type=str)
    adv_parser.add_argument('-Dp', '--discriminator-load-path', type=str)
    adv_parser.add_argument('-Go', '--generator-optimizer', type=str, default='Adam')
    adv_parser.add_argument('-Do', '--discriminator-optimizer', type=str, default='RMSprop')
    adv_parser.add_argument('-Gg', '--generator-gradient-clipping', type=float, default=False)
    adv_parser.add_argument('-Dg', '--discriminator-gradient-clipping', type=float, default=False)
    adv_parser.add_argument('-Gl', '--generator-learning-rate', type=float, default=0.001)
    adv_parser.add_argument('-Dl', '--discriminator-learning-rate', type=float, default=0.001)
    adv_parser.add_argument('-Gr', '--reset-generator-optimizer', action='store_true', default=False)
    adv_parser.add_argument('-Dr', '--reset-discriminator-optimizer', action='store_true', default=False)
    adv_parser.add_argument('-Gq', '--g-max-q-mean', type=float, default=0.5)
    adv_parser.add_argument('-Gm', '--g-min-q-mean', type=float, default=0.0)
    adv_parser.add_argument('-Dm', '--d-min-loss', type=float, default=0.5)
    adv_parser.add_argument('-Gs', '--g-max-steps', type=int, default=0)
    adv_parser.add_argument('-Ds', '--d-max-steps', type=int, default=0)
    adv_parser.add_argument('-f', '--first-to-train', type=str, default='G', choices=['G', 'D'])
    adv_parser.add_argument('-ms', '--monte-carlo-sample-size', type=int, default=8)
    adv_parser.add_argument('-mf', '--monte-carlo-sample-factor', type=float, default=1.0)
    adv_parser.add_argument('-c', '--control-ratio', type=float, default=1.0)
    adv_parser.add_argument('-i', '--save-interval', type=float, default=60.0)
    adv_parser.add_argument('-L', '--enable-logging', action='store_true', default=False)
    adv_parser.set_defaults(main=adversarial)
    return parser.parse_args()

if __name__ == '__main__':
    args = get_args()
    args.main(args)


""" Proof of concept for n-way matching using footprint detection

Some terminology:

cell    :   A small area, used to find matches,  The size
    should be about the same as the maximum match radius.

source   :  As per DM, per-catalog detection

matchWcs :  The WCS used to define the cells

subRegion : A square sub-region of the skymap defined by the matchWcs

sourceCountsMap :  A map of the number of source per cell, made per
   subRegion

Cluster:  A set of sources found by running a footprint finding algorithm
   on a sourceCountsMap

"""

import sys
import os
import glob

from collections import OrderedDict

import time
import numpy as np
from astropy import wcs
from astropy.table import Table
from astropy.table import vstack
from astropy.io import fits

try:
    import pyarrow.parquet as pq  # noqa
except ImportError:
    print("nway requires pyarrrow")

try:
    import lsst.afw.detection as afwDetect
    import lsst.afw.image as afwImage
except ImportError:
    print("nway requires lsst.afw")

RECURSE_MAX = 200
COLUMNS = ['ra', 'decl', 'visit', 'ccd', 'sky_source', 'sourceId', 'PsFlux', 'PsFluxErr', 'Centroid_flag', 'detect_isPrimary']

def createGlobalWcs(refDir, cellSize, nCell):
    """ Helper function to create the WCS used to project the
    sources in a skymap """
    w = wcs.WCS(naxis=2)
    w.wcs.cdelt = [-cellSize, cellSize]
    w.wcs.crpix = [nCell[0]/2, nCell[1]/2]
    w.wcs.crval = [refDir[0], refDir[1]]
    return w

def clusterStats(clusterDict):
    """ Helper function to get stats about the clusters

    'Orphan'   means single source clusters (i.e., single detections)
    'Mixed`    means there is more that one source from at least one
               input catalog
    'Confused' means there are more than four cases of duplication
    """
    nOrphan = 0
    nMixed = 0
    nConfused = 0
    for val in clusterDict.values():
        if val.nSrc == 1:
            nOrphan += 1
        if val.nSrc != val.nUnique:
            nMixed += 1
            if val.nSrc > val.nUnique + 3:
                nConfused += 1
    return np.array([len(clusterDict), nOrphan, nMixed, nConfused])

class ClusterData:
    """ Class to store data about clusters

    Parameters
    ----------
    iCluster : `int`
        Cluster ID
    origCluster : `int`
        Id of the original cluster this cluster was made from
    nSrc : `int`
        Number of sources in this cluster
    nUnique : `int`
        Number of catalogs contributing sources to this cluster
    catIndices : `np.array`, [`int`]
        Indices of the catalogs of sources associated to this cluster
    sourcdIds : `np.array`, [`int`]
        Sources IDs of the sources associated to this cluster
    sourcdIdxs : `np.array`, [`int`]
        Indices of the sources with their respective catalogs
    xCent : `float`
        X-pixel value of cluster centroid (in WCS used to do matching)
    yCent : `float`
        Y-pixel value of cluster centroid (in WCS used to do matching)
    """
    def __init__(self, iCluster, footprint, sources, origCluster=None):
        self._iCluster = iCluster
        self._footprint = footprint
        if origCluster is None:
            self._origCluster = self._iCluster
        else:
            self._origCluster = origCluster
        self._catIndices = sources[0]
        self._sourceIds = sources[1]
        self._sourceIdxs = sources[2]
        self._nSrc =  self._catIndices.size
        self._nUnique = len(np.unique(self._catIndices))
        self._objects = []
        self._xCent = None
        self._yCent = None
        self._dist2 = None
        self._rmsDist = None
        self.xCell = None
        self.yCell = None
        self.snr = None

    def extract(self, subRegionData):
        """ Extract the xCell, yCell and snr data from
        the sources in this cluster
        """
        self.xCell = np.zeros((self._nSrc), np.float32)
        self.yCell = np.zeros((self._nSrc), np.float32)
        self.snr = np.zeros((self._nSrc), np.float32)
        for i, (iCat, srcIdx) in enumerate(zip(self._catIndices, self._sourceIdxs)):
            self.xCell[i] = subRegionData.data[iCat]['xcell'].values[srcIdx]
            self.yCell[i] = subRegionData.data[iCat]['ycell'].values[srcIdx]
            self.snr[i] = subRegionData.data[iCat]['SNR'].values[srcIdx]

    def clearTempData(self):
        """ Remove temporary data only used when making objects """
        self.xCell = None
        self.yCell = None
        self.snr = None

    @property
    def iCluster(self):
        """ Return the cluster ID """
        return self._iCluster

    @property
    def nSrc(self):
        """ Return the number of sources associated to the cluster """
        return self._nSrc

    @property
    def nUnique(self):
        """ Return the number of catalogs contributing sources to the cluster """
        return self._nUnique

    @property
    def sourceIds(self):
        """ Return the source IDs associated to this cluster """
        return self._sourceIds

    @property
    def dist2(self):
        """ Return an array with the distance squared (in cells)
        between each source and the cluster centroid """
        return self._dist2

    @property
    def objects(self):
        """ Return the objects associated with this cluster """
        return self._objects

    def processCluster(self, subRegionData, pixelR2Cut):
        """ Function that is called recursively to
        split clusters until they:

        1.  Consist only of sources with the match radius of the cluster
        centroid.

        2.  Have at most one source per input catalog
        """
        self._nSrc =  self._catIndices.size
        self._nUnique = len(np.unique(self._catIndices))
        if self._nSrc == 0:
            print("Empty cluster", self._nSrc, self._nUnique)
            return self._objects
        self.extract(subRegionData)
        if self._nSrc == 1:
            self._xCent = self.xCell[0]
            self._yCent = self.yCell[0]
            self._dist2 = np.zeros((1))
            self._rmsDist = 0.
            initialObject = self.addObject(subRegionData)
            initialObject.processObject(subRegionData, pixelR2Cut)
            self.clearTempData()
            return self._objects

        sumSnr = np.sum(self.snr)
        self._xCent = np.sum(self.xCell*self.snr) / sumSnr
        self._yCent = np.sum(self.yCell*self.snr) / sumSnr
        self._dist2 = (self._xCent - self.xCell)**2 + (self._yCent - self.yCell)**2
        self._rmsDist = np.sqrt(np.mean(self._dist2))

        initialObject = self.addObject(subRegionData)
        initialObject.processObject(subRegionData, pixelR2Cut)
        self.clearTempData()
        return self._objects

    def addObject(self, subRegionData, mask=None):
        """ Add a new object to this cluster """
        newObject = subRegionData.addObject(self, mask)
        self._objects.append(newObject)
        return newObject

class ObjectData:
    """ Small class to define 'Objects', i.e., sets of associated sources """

    def __init__(self, cluster, objectId, mask):
        """ Build from `ClusterData`, an objectId and mask specifying with sources
        in the cluster are part of the object """
        self._parentCluster = cluster
        self._objectId = objectId
        if mask is None:
            self._mask = np.ones((self._parentCluster.nSrc), dtype=bool)
        else:
            self._mask = mask
        self._catIndices = self._parentCluster._catIndices[self._mask]
        self._nSrc = self._catIndices.size
        self._nUnique = np.unique(self._catIndices).size
        self._xCent = None
        self._yCent = None
        self._dist2 = None
        self._rmsDist = None

    @property
    def nSrc(self):
        """ Return the number of sources associated to the cluster """
        return self._nSrc

    @property
    def nUnique(self):
        """ Return the number of catalogs contributing sources to the cluster """
        return self._nUnique

    @property
    def dist2(self):
        """ Return an array with the distance squared (in cells)
        between each source and the cluster centroid """
        return self._dist2

    def updateCatIndices(self):
        self._catIndices = self._parentCluster._catIndices[self._mask]
        self._nSrc = self._catIndices.size
        self._nUnique = np.unique(self._catIndices).size

    def sourceIds(self):
        return self._parentCluster.sourceIds[self._mask]

    def processObject(self, subRegionData, pixelR2Cut, recurse=0):
        """ Recursively process an object and make sub-objects """
        if recurse > RECURSE_MAX:
            print("Recursion limit: ", self._nSrc, self._nUnique)
            return
        if self._nSrc == 0:
            print("Empty object", self._nSrc, self._nUnique, recurse)
            return

        xCell = self._parentCluster.xCell[self._mask]
        yCell = self._parentCluster.yCell[self._mask]
        snr = self._parentCluster.snr[self._mask]

        if self._mask.sum() == 1:
            self._xCent = xCell[0]
            self._yCent = yCell[0]
            self._dist2 = np.zeros((1), float)
            self._rmsDist = 0.
            return

        sumSnr = np.sum(snr)
        self._xCent = np.sum(xCell*snr) / sumSnr
        self._yCent = np.sum(yCell*snr) / sumSnr
        self._dist2 = np.array((self._xCent - xCell)**2 + (self._yCent - yCell)**2)
        self._rmsDist = np.sqrt(np.mean(self._dist2))
        subMask = self._dist2 < pixelR2Cut
        if subMask.all():
            if self._nSrc != self._nUnique:
                self.splitObject(subRegionData, pixelR2Cut, recurse=recurse+1)
            return

        if not subMask.any():
            idx = np.argmax(snr)
            self._xCent = xCell[idx]
            self._yCent = yCell[idx]
            self._dist2 = np.array((self._xCent - xCell)**2 + (self._yCent - yCell)**2)
            self._rmsDist = np.sqrt(np.mean(self._dist2))
            subMask = self._dist2 < pixelR2Cut

        newObjMask = self._mask.copy()
        newObjMask[newObjMask] *= subMask

        newObject = self._parentCluster.addObject(subRegionData, newObjMask)
        newObject.processObject(subRegionData, pixelR2Cut)

        self._mask[self._mask] *= ~subMask
        self.updateCatIndices()
        self.processObject(subRegionData, pixelR2Cut, recurse=recurse+1)

    def splitObject(self, subRegionData, pixelR2Cut, recurse=0):
        """ Split up a cluster keeping only one source per input
        catalog, choosing the one closest to the cluster center """
        sortIdx = np.argsort(self._dist2)
        mask = np.ones((self._nSrc), dtype=bool)
        usedCats = {}
        for iSrc, catIdx in zip(sortIdx, self._catIndices[sortIdx]):
            if catIdx not in usedCats:
                usedCats[catIdx] = 1
                continue
            else:
                usedCats[catIdx] += 1
            mask[iSrc] = False

        newObjMask = self._mask.copy()
        newObjMask[newObjMask] *= mask

        newObject = self._parentCluster.addObject(subRegionData, newObjMask)
        newObject.processObject(subRegionData, pixelR2Cut)

        self._mask[self._mask] *= ~mask
        self.updateCatIndices()
        self.processObject(subRegionData, pixelR2Cut, recurse=recurse+1)

class SubregionData:
    """ Class to analyze data for a SubRegion

    Include sub-region boundries, reduced data tables
    and clustering results

    Does not store sky maps

    Subregions are square sub-regions of the Skymap
    constructed with the WCS

    The subregion covers corner:corner+size

    The sources are projected into an array that extends `buf` cells
    beyond the region.

    Parameters
    ----------
    _data : `list`, [`Dataframe`]
        Reduced dataframes with only sources for this sub-region

    _clusterIds : `list`, [`np.array`]
        Matched arrays with the index of the cluster associated to each
        source.  I.e., these could added to the Dataframes as
        additional columns

    _clusterDict : `dict`, [`int` : `ClusterData`]
        Dictionary with cluster membership data

    TODO:  Add code to filter out clusters centered in the buffer
    """
    def __init__(self, matcher, idOffset, corner, size, buf=10):
        self._matcher = matcher
        self._idOffset = idOffset # Offset used for the Object and Cluster IDs for this region
        self._corner = corner # cellX, cellY for corner of region
        self._size = size # size of region
        self._buf = buf
        self._minCell = corner - buf
        self._maxCell = corner + size + buf
        self._nCells = self._maxCell - self._minCell
        self._data = None
        self._nSrc = None
        self._footprintIds = None
        self._clusterDict = OrderedDict()
        self._objectDict = OrderedDict()

    def reduceData(self, data):
        """ Pull out only the data needed for this sub-region """
        self._data = [self.reduceDataframe(val) for val in data]
        self._nSrc = sum([len(df) for df in self._data])

    @property
    def nClusters(self):
        """ Return the number of clusters in this region """
        return len(self._clusterDict)

    @property
    def nObjects(self):
        """ Return the number of objects in this region """
        return len(self._objectDict)

    @property
    def data(self):
        """ Return the data associated to this region """
        return self._data

    @property
    def clusterDist(self):
        """ Return a dictionary mapping clusters Ids to clusters """
        return self._clusterDict

    def reduceDataframe(self, dataframe):
        """ Filters dataframe to keep only source in the subregion """
        xLocal = dataframe['xcell'] - self._minCell[0]
        yLocal = dataframe['ycell'] - self._minCell[1]
        filtered = (xLocal >= 0) & (xLocal < self._nCells[0]) & (yLocal >= 0) & (yLocal < self._nCells[1])
        red = dataframe[filtered].copy(deep=True)
        red['xlocal'] = xLocal[filtered]
        red['ylocal'] = yLocal[filtered]
        return red

    def countsMap(self, weightName=None):
        """ Fill a map that counts the number of source per cell """
        toFill = np.zeros((self._nCells))
        for df in self._data:
            toFill += self.fillSubRegionFromDf(df, weightName=weightName)
        return toFill

    def associateSourcesToFootprints(self, clusterKey):
        """ Loop through data and associate sources to clusters """
        self._footprintIds = [self.findClusterIds(df, clusterKey) for df in self._data]

    def buildClusterData(self, fpSet, pixelR2Cut=4.):
        """ Loop through cluster ids and collect sources into
        the ClusterData objects """
        footprints = fpSet.getFootprints()
        footprintDict = {}
        nMissing = 0
        nFound = 0
        for iCat, (df, footprintIds) in enumerate(zip(self._data, self._footprintIds)):
            for srcIdx, (srcId, footprintId) in enumerate(zip(df['sourceId'], footprintIds)):
                if footprintId < 0:
                    nMissing += 1
                    continue
                if footprintId not in footprintDict:
                    footprintDict[footprintId] = [(iCat, srcId, srcIdx)]
                else:
                    footprintDict[footprintId].append((iCat, srcId, srcIdx))
                nFound += 1
        for footprintId, sources in footprintDict.items():
            footprint = footprints[footprintId]
            iCluster = footprintId+self._idOffset
            cluster = ClusterData(iCluster, footprint, np.array(sources).T)
            self._clusterDict[iCluster] = cluster
            cluster.processCluster(self, pixelR2Cut)

    def analyze(self, weightName=None, pixelR2Cut=4.):
        """ Analyze this sub-region

        Note that this returns the counts maps and clustering info,
        which can be helpful for debugging.
        """
        if self._nSrc == 0:
            return None
        countsMap = self.countsMap(weightName)
        oDict = self.getFootprints(countsMap)
        oDict['countsMap'] = countsMap
        self.associateSourcesToFootprints(oDict['footprintKey'])
        self.buildClusterData(oDict['footprints'], pixelR2Cut)
        return oDict

    @staticmethod
    def findClusterIds(df, clusterKey):
        """ Associate sources to clusters using `clusterkey`
        which is a map where any pixel associated to a cluster
        has the cluster index as its value """
        return np.array([clusterKey[yLocal,xLocal] for xLocal, yLocal in zip(df['xlocal'], df['ylocal'])]).astype(np.int32)

    def fillSubRegionFromDf(self, df, weightName=None):
        """ Fill a source counts map from a reduced dataframe for one input
        catalog """
        if weightName is None:
            weights = None
        else:
            weights = df[weightName].values
        hist = np.histogram2d(df['xlocal'], df['ylocal'], bins=self._nCells,
                              range=((0, self._nCells[0]),
                                     (0, self._nCells[1])),
                              weights=weights)
        return hist[0]

    @staticmethod
    def filterFootprints(fpSet, buf):
        """ Remove footprints within `buf` cells of the region edge """
        region = fpSet.getRegion()
        width, height = region.getWidth(), region.getHeight()
        outList = []
        maxX = width - buf
        maxY = height - buf
        for fp in fpSet.getFootprints():
            cent = fp.getCentroid()
            xC = cent.getX()
            yC = cent.getY()
            if xC < buf or xC > maxX or yC < buf or yC > maxY:
                continue
            outList.append(fp)
        fpSetOut = afwDetect.FootprintSet(fpSet.getRegion())
        fpSetOut.setFootprints(outList)
        return fpSetOut

    def getFootprints(self, countsMap):
        """ Take a source counts map and do clustering using Footprint detection
        """
        image = afwImage.ImageF(countsMap.astype(np.float32))
        footprintsOrig = afwDetect.FootprintSet(image, afwDetect.Threshold(0.5))
        footprints = self.filterFootprints(footprintsOrig, self._buf)
        footprintKey = afwImage.ImageI(np.full(countsMap.shape, -1, dtype=np.int32))
        for i, footprint in enumerate(footprints.getFootprints()):
            footprint.spans.setImage(footprintKey, i, doClip=True)
        return dict(image=image, footprints=footprints, footprintKey=footprintKey)

    def getClusterAssociations(self):
        """ Convert the clusters to a set of associations """
        clusterIds = []
        sourceIds = []
        distances = []
        for cluster in self._clusterDict.values():
            clusterIds.append(np.full((cluster.nSrc), cluster.iCluster, dtype=int))
            sourceIds.append(cluster.sourceIds)
            distances.append(cluster.dist2)
        if not distances:
            return Table(dict(distance=[], id=np.array([], int), object=np.array([], int)))
        distances = np.hstack(distances)
        distances = self._matcher.cellToArcsec() * np.sqrt(distances)
        data = dict(object=np.hstack(clusterIds),
                    id=np.hstack(sourceIds),
                    distance=distances)
        return Table(data)

    def getObjectAssociations(self):
        clusterIds = []
        objectIds = []
        sourceIds = []
        distances = []
        for obj in self._objectDict.values():
            clusterIds.append(np.full((obj._nSrc), obj._parentCluster.iCluster, dtype=int))
            objectIds.append(np.full((obj._nSrc), obj._objectId, dtype=int))
            sourceIds.append(obj.sourceIds())
            distances.append(obj.dist2)
        if not distances:
            return Table(dict(object=np.array([], int),
                              parent=np.array([], int),
                              id=np.array([], int),
                              distance=[]))
        distances = np.hstack(distances)
        distances = self._matcher.cellToArcsec() * np.sqrt(distances)
        data = dict(object=np.hstack(objectIds),
                    parent=np.hstack(clusterIds),
                    id=np.hstack(sourceIds),
                    distance=distances)
        return Table(data)

    def getClusterStats(self):
        """ Convert the clusters to a set of associations """
        nClust = self.nClusters
        clusterIds = np.zeros((nClust), dtype=int)
        nSrcs = np.zeros((nClust), dtype=int)
        nObjects = np.zeros((nClust), dtype=int)
        nUniques = np.zeros((nClust), dtype=int)
        distRms = np.zeros((nClust), dtype=float)
        xCents = np.zeros((nClust), dtype=float)
        yCents = np.zeros((nClust), dtype=float)
        for idx, cluster in enumerate(self._clusterDict.values()):
            clusterIds[idx] = cluster._iCluster
            nSrcs[idx] = cluster.nSrc
            nObjects[idx] = len(cluster._objects)
            nUniques[idx] = cluster.nUnique
            distRms[idx] = cluster._rmsDist
            xCents[idx] = cluster._xCent
            yCents[idx] = cluster._yCent
        ra, decl = self._matcher.cellToWorld(xCents, yCents)
        distRms *= self._matcher.cellToArcsec()

        data = dict(clusterIds=clusterIds,
                    nSrcs=nSrcs,
                    nObject=nObjects,
                    nUnique=nUniques,
                    distRms=distRms,
                    ra=ra,
                    decl=decl)

        return Table(data)

    def getObjectStats(self):
        """ Convert the clusters to a set of associations """
        nObj = self.nObjects
        clusterIds = np.zeros((nObj), dtype=int)
        objectIds = np.zeros((nObj), dtype=int)
        nSrcs = np.zeros((nObj), dtype=int)
        distRms = np.zeros((nObj), dtype=float)
        xCents = np.zeros((nObj), dtype=float)
        yCents = np.zeros((nObj), dtype=float)
        for idx, obj in enumerate(self._objectDict.values()):
            clusterIds[idx] = obj._parentCluster._iCluster
            objectIds[idx] = obj._objectId
            nSrcs[idx] = obj.nSrc
            distRms[idx] = obj._rmsDist
            xCents[idx] = obj._xCent
            yCents[idx] = obj._yCent

        ra, decl = self._matcher.cellToWorld(xCents, yCents)
        distRms *= self._matcher.cellToArcsec()

        data = dict(clusterIds=clusterIds,
                    objectIds=objectIds,
                    nSrcs=nSrcs,
                    distRms=distRms,
                    ra=ra,
                    decl=decl)

        return Table(data)

    def addObject(self, cluster, mask=None):
        """ Add an object to this sub-region """
        objectId = self.nObjects + self._idOffset
        newObject = ObjectData(cluster, objectId, mask)
        self._objectDict[objectId] = newObject
        return newObject

class NWayMatch:
    """ Class to do N-way matching

    Uses a provided WCS to define a Skymap that covers the full region
    begin matched.

    Uses that WCS to assign cell locations to all sources in the input catalogs

    Iterates over sub-regions and does source clustering in each sub-region
    using Footprint detection on a Skymap of source counts per cell.

    Assigns each input source to a cluster.

    At that stage the clusters are not the final product as they can include
    more than one soruce from a given catalog.

    Loops over clusters and processes each cluster to:

       1. Remove outliers outside the match radius w.r.t. the cluster centroid.
       2. Resolve cases of confusion, where multiple sources from a single
       catalog contribute to a cluster.

    Parameters
    ----------
    _redData : `list`, [`Dataframe`]
        Reduced dataframes with only the columns needed for matching

    _clusters : `OrderedDict`, [`tuple`, `SubregionData`]
        Dictionary providing access to subregion data
    """

    def __init__(self, matchWcs, **kwargs):
        self._wcs = matchWcs
        self._cellSize = self._wcs.wcs.cdelt[1]
        self._nCellSide = np.ceil(2*np.array(self._wcs.wcs.crpix)).astype(int)
        self._subRegionSize = kwargs.get('subRegionSize', 3000)
        self._subRegionBuffer = kwargs.get('subRegionBuffer', 10)
        self._subregionMaxObject = kwargs.get('subregionMaxObject', 100000)
        self._pixelR2Cut = kwargs.get('pixelR2Cut', 1.0)
        self._nSubRegion = np.ceil(self._nCellSide/self._subRegionSize)
        self._redData = OrderedDict()
        self._clusters = None

    def cellToArcsec(self):
        return 3600. * self._cellSize

    def cellToWorld(self, xCell, yCell):
        return self._wcs.wcs_pix2world(xCell, yCell, 0)

    @classmethod
    def create(cls, refDir, regionSize, cellSize, **kwargs):
        """ Make an `NWayMatch` object from inputs """
        nCell = (np.array(regionSize)/cellSize).astype(int)
        matchWcs = createGlobalWcs(refDir, cellSize, nCell)
        return cls(matchWcs, **kwargs)

    @property
    def redData(self):
        """ Return the dictionary of reduced data, i.e., just the columns
        need for matching """
        return self._redData

    @property
    def nSubRegion(self):
        """ Return the number of sub-regions in X,Y """
        return self._nSubRegion

    def reduceData(self, inputFiles, visitIds):
        """ Read input files and filter out only the columns we need """
        for fName, vid in zip(inputFiles, visitIds):
            self._redData[vid] = self.reduceDataFrame(fName)

    def reduceDataFrame(self, fName):
        """ Read and reduce a single input file """
        parq = pq.read_pandas(fName, columns=COLUMNS)
        df = parq.to_pandas()
        df['SNR'] = df['PsFlux']/df['PsFluxErr']
        # select sources that have SNR > 5.
        # You may start with 10 or even 50 if you want to start with just the brightest objects
        # AND
        # Centroid_flag is True if there was a problem fitting the position (centroid)
        # AND
        # sky_source is True if it is a measurement of blank sky.
        # sky_sources should have SNR < 5 or the Centroid_flag set,
        # but explicitly filter just to make sure.
        # AND
        # detect_isPrimary = True to remove duplicate rows from deblending:
        # If a source has been deblended, the parent is marked detect_isPrimary=False and its children True.
        df_clean = df[(df.SNR > 5) & ~df.Centroid_flag & ~df.sky_source & df.detect_isPrimary]
        xcell, ycell = self._wcs.wcs_world2pix(df_clean['ra'].values, df_clean['decl'].values, 0)
        df_red = df_clean[["ra", "decl", "SNR", "sourceId"]].copy(deep=True)
        df_red['xcell'] = xcell
        df_red['ycell'] = ycell
        return df_red[["ra", "decl", "SNR", "sourceId", "xcell", "ycell"]]

    def reduceCatalog(self, catalog):
        """ Reduce a catalog """
        raise NotImplementedError()

    def add(self, catalog, vid):
        """ Add a catalog to the data set being matched """
        self._redData[vid] = self.reduceCatalog(catalog)

    def getIdOffset(self, ix, iy):
        """ Get the ID offset to use for a given sub-region """
        subRegionIdx = self._nSubRegion[1]*ix + iy
        return int(self._subregionMaxObject * subRegionIdx)

    def analyzeSubregion(self, ix, iy, fullData=False):
        """ Analyze a single subregion

        Returns an OrderedDict

        'srd' : `SubregionData`
            The analysis data for the sub-region

        if fullData is True the return dict will include

        'image' : `afwImage.ImageI`
            Image of subregion source counts map
        'countsMap' : `np.array`
            Numpy array with same
        'clusters' : `afwDetect.FootprintSet`
            Clusters as dectected by finding FootprintSet on source counts map
        'clusterKey' : `afwImage.ImageI`
            Map of subregion with pixels filled with index of
            associated Footprints
        """
        iSubRegion = np.array([ix, iy])
        corner = iSubRegion * self._subRegionSize
        idOffset = self.getIdOffset(ix, iy)
        srd = SubregionData(self, idOffset, corner, self._subRegionSize, self._subRegionBuffer)
        srd.reduceData(self._redData.values())
        oDict = srd.analyze(pixelR2Cut=self._pixelR2Cut)
        if oDict is None:
            return None
        if fullData:
            oDict['srd'] = srd
            return oDict
        if srd.nObjects >= self._subregionMaxObject:
            print("Too many object in a subregion", srd.nObjects, elf._subregionMaxObject)
        return dict(srd=srd)

    def finish(self):
        """ Does clusering for all subregions

        Does not store source counts maps for the counts regions
        """
        self._clusters = OrderedDict()
        nAssoc = 0
        clusterAssocTables = []
        objectAssocTables = []
        clusterStatsTables = []
        objectStatsTables = []

        for ix in range(int(self._nSubRegion[0])):
            sys.stdout.write("%2i " % ix)
            sys.stdout.flush()
            for iy in range(int(self._nSubRegion[1])):
                sys.stdout.write('.')
                sys.stdout.flush()
                iSubRegion = (ix, iy)
                odict = self.analyzeSubregion(ix, iy)
                if odict is None:
                    continue
                subregionData = odict['srd']
                self._clusters[iSubRegion] = subregionData
                clusterAssocTables.append(subregionData.getClusterAssociations())
                objectAssocTables.append(subregionData.getObjectAssociations())
                clusterStatsTables.append(subregionData.getClusterStats())
                objectStatsTables.append(subregionData.getObjectStats())

            sys.stdout.write('!\n')

        sys.stdout.write("Making association vectors\n")
        hduList = fits.HDUList([fits.PrimaryHDU(),
                                fits.table_to_hdu(vstack(clusterAssocTables)),
                                fits.table_to_hdu(vstack(objectAssocTables)),
                                fits.table_to_hdu(vstack(clusterStatsTables)),
                                fits.table_to_hdu(vstack(objectStatsTables))])
        return hduList

    def allStats(self):
        """ Helper function to print info about clusters """
        stats = np.zeros((4), int)
        for key, srd in self._clusters.items():
            subRegionStats = clusterStats(srd._clusterDict)
            print("%3i, %3i: %8i %8i %8i %8i" % (key[0], key[1], subRegionStats[0], subRegionStats[1], subRegionStats[2], subRegionStats[3]))
            stats += subRegionStats
        return stats

def main():
    """ Example usage """

    DATADIR = "."
    SOURCE_TABLEFILES = glob.glob(os.path.join(DATADIR, "sourceTable-*.parq"))
    VISIT_IDS = np.arange(len(SOURCE_TABLEFILES))

    REF_DIR = (150., 2.)  # RA, DEC in deg
    REGION_SIZE = (3., 3.)  # in Deg
    #CELL_SIZE = 5.0e-5    # in Deg
    CELL_SIZE = 1. / (3600*2) # in Deg
    #SUBREGION_SIZE = 2700 # in Pixels
    SUBREGION_SIZE = 1350 # in Pixels
    PIXEL_R2CUT = 1.

    t0 = time.time()
    nWay = NWayMatch.create(REF_DIR, REGION_SIZE, CELL_SIZE, pixelR2Cut=PIXEL_R2CUT, subRegionSize=SUBREGION_SIZE)
    print("Building clusters in %ix%i sub-regions" % (nWay.nSubRegion[0], nWay.nSubRegion[1]))
    nWay.reduceData(SOURCE_TABLEFILES, VISIT_IDS)
    outTables = nWay.finish()
    t1 = time.time()
    print("Reading and clustering took %s s" % (t1-t0))

    print("Cluster Summaries for sub-regions")
    print("Region  :  nCluster nOrphan  nMixed   nConf")
    stats = nWay.allStats()
    print("Total:   %8i %8i %8i %8i" % (stats[0], stats[1], stats[2], stats[3]))

    outTables.writeto("out.fits", overwrite=True)

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
Spectroscope - Analysis
=======================

Defines the objects for the homemade spectroscope spectrum images analysis.

References
----------
.. [1]  http://thomasmansencal.blogspot.fr/2014/07/a-homemade-spectroscope.html
"""

from __future__ import division, unicode_literals

import numpy as np
import scipy.ndimage

from colour import (Extrapolator, LinearInterpolator, RGB_COLOURSPACES,
                    RGB_luminance, SpectralDistribution,
                    MultiSpectralDistributions)
from colour.utilities import tstack

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2021 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'RGB_Spectrum', 'image_profile', 'calibrate_RGB_spectrum_profile',
    'calibrated_RGB_spectrum', 'luminance_sd'
]

class RGB_Spectrum(MultiSpectralDistributions):
    """
    Defines an *RGB* spectrum object implementation.

    Parameters
    ----------
    data : Series or Dataframe or Signal or MultiSignals or \
MultiSpectralDistributions or array_like or dict_like, optional
        Data to be stored in the multi-spectral distributions.
    domain : array_like, optional
        Values to initialise the multiple :class:`colour.SpectralDistribution`
        class instances :attr:`colour.continuous.Signal.wavelengths` attribute
        with. If both ``data`` and ``domain`` arguments are defined, the latter
        will be used to initialise the
        :attr:`colour.continuous.Signal.wavelengths` attribute.
    labels : array_like, optional
        Names to use for the :class:`colour.SpectralDistribution` class
        instances.

    Other Parameters
    ----------------
    name : unicode, optional
       Multi-spectral distribution name.
    interpolator : object, optional
        Interpolator class type to use as interpolating function for the
        :class:`colour.SpectralDistribution` class instances.
    interpolator_args : dict_like, optional
        Arguments to use when instantiating the interpolating function
        of the :class:`colour.SpectralDistribution` class instances.
    extrapolator : object, optional
        Extrapolator class type to use as extrapolating function for the
        :class:`colour.SpectralDistribution` class instances.
    extrapolator_args : dict_like, optional
        Arguments to use when instantiating the extrapolating function
        of the :class:`colour.SpectralDistribution` class instances.
    strict_labels : array_like, optional
        Multi-spectral distribution labels for figures, default to
        :attr:`colour.characterisation.RGB_SpectralSensitivities.labels`
        attribute value.
    """

    def __init__(self, data=None, domain=None, labels=None, **kwargs):
        super(RGB_Spectrum, self).__init__(
            data, domain, labels=('R', 'G', 'B'), **kwargs)

def image_profile(image, line, samples=None):
    """
    Returns the image profile using given line coordinates and given samples
    count.

    Parameters
    ----------
    image : ndarray
        Image to retrieve the profile.
    line : tuple or list or ndarray, (x0, y0, x1, y1)
        Coordinates as image array indexes to measure the profile.
    samples : int, optional
        Samples count to retrieve along the line, default to image width.

    Returns
    -------
    ndarray
        Profile.

    References
    ----------
    .. [2]  http://stackoverflow.com/a/7880726/931625
            (Last accessed 8 August 2014)
    """

    height, width, channels = image.shape
    samples = samples if samples else width
    x0, y0, x1, y1 = line

    profile = []
    for i in range(channels):
        x, y = np.linspace(x0, x1, samples), np.linspace(y0, y1, samples)
        z = image[:, :, i]

        profile.append(
            scipy.ndimage.map_coordinates(np.transpose(z), np.vstack([x, y])))

    return np.dstack(profile)

def calibrate_RGB_spectrum_profile(profile, reference, measured, samples=None):
    """
    Calibrates given spectrum profile using given theoretical reference
    wavelength lines in nanometers and measured lines in horizontal axis pixels
    values. If more than 2 lines are provided the profile data will be warped
    to fit the theoretical reference wavelength lines.

    Parameters
    ----------
    profile : ndarray
        Image profile to calibrate.
    reference : dict
        Theoretical reference wavelength lines.
    measured : dict
        Measured lines in horizontal axis pixels values.
    samples : int, optional
        Profile samples count.

    Returns
    -------
    RGB_Spectrum
        Calibrated RGB spectrum.
    """

    samples = samples if samples else profile.shape[1]
    measured_lines = [
        line for line, value in sorted(measured.items(), key=lambda x: x[1])
    ]

    # Reference samples.
    r = np.array([reference.get(sample) for sample in measured_lines])
    # Measured samples.
    m = np.array([measured.get(sample) for sample in measured_lines])

    # Reference range array.
    rr = np.linspace(min(r), max(r))
    # Measured range array.
    mm = np.linspace(min(m), max(m))

    # Interpolator from reference to measured.
    r_to_m_interpolator = Extrapolator(LinearInterpolator(r, m))

    # Interpolator from measured range to reference range.
    mm_to_rr_interpolator = Extrapolator(LinearInterpolator(mm, rr))

    # Colors interpolator.
    R_interpolator = Extrapolator(
        LinearInterpolator(np.arange(0, profile.shape[1]), profile[0, :, 0]))
    G_interpolator = Extrapolator(
        LinearInterpolator(np.arange(0, profile.shape[1]), profile[0, :, 1]))
    B_interpolator = Extrapolator(
        LinearInterpolator(np.arange(0, profile.shape[1]), profile[0, :, 2]))

    wavelengths = np.linspace(
        mm_to_rr_interpolator([0]), mm_to_rr_interpolator([profile.shape[1]]),
        samples)

    return RGB_Spectrum(
        dict(
            zip(wavelengths,
                tstack([
                    R_interpolator(r_to_m_interpolator(wavelengths)),
                    G_interpolator(r_to_m_interpolator(wavelengths)),
                    B_interpolator(r_to_m_interpolator(wavelengths))
                ]))),
        name='RGB Spectrum')

def calibrated_RGB_spectrum(image, reference, measured, samples=None):
    """
    Returns the RGB spectrum of given image.

    Parameters
    ----------
    image : ndarray
        Image to retrieve the RGB spectrum, assuming the spectrum is already
        properly oriented.
    reference : dict
        Theoretical reference wavelength lines.
    measured : dict
        Measured lines in horizontal axis pixels values.
    samples : int, optional
        Spectrum samples count.

    Returns
    -------
    RGB_Spectrum
        RGB spectrum.
    """

    samples = samples if samples else image.shape[1]
    profile = image_profile(
        image, line=[0, 0, image.shape[1] - 1, 0], samples=samples)

    return calibrate_RGB_spectrum_profile(
        profile=profile,
        reference=reference,
        measured=measured,
        samples=samples)

def luminance_sd(spectrum, colourspace=RGB_COLOURSPACES['sRGB']):
    """
    Returns the luminance spectral distribution of given RGB spectrum.

    Parameters
    ----------
    spectrum : RGB_Spectrum
        RGB spectrum to retrieve the luminance from.
    colourspace : RGB_Colourspace
        *RGB* Colourspace.

    Returns
    -------
    SpectralDistribution
        RGB spectrum luminance spectral distribution, units are arbitrary
        and normalised to [0, 100] domain.
    """

    spectrum = spectrum.copy().normalise(100)
    luminance = lambda x: RGB_luminance(x, colourspace.primaries, colourspace.
                                        whitepoint)

    return SpectralDistribution(
        dict(zip(spectrum.wavelengths, luminance(spectrum.values))),
        name='calibrated_RGB_spectrum')


"""
Classes for specific source models
"""
import numpy as np
from yt.funcs import ensure_numpy_array
from tqdm.auto import tqdm
from pyxsim.utils import mylog
from yt.units.yt_array import YTQuantity
from yt.utilities.physical_constants import clight
from pyxsim.spectral_models import thermal_models
from pyxsim.utils import parse_value, isunitful
from soxs.utils import parse_prng
from soxs.constants import elem_names, atomic_weights, metal_elem
from yt.utilities.exceptions import YTFieldNotFound
from yt.utilities.parallel_tools.parallel_analysis_interface import \
    parallel_objects, communication_system, parallel_capable
from numbers import Number

comm = communication_system.communicators[-1]

solar_H_abund = 0.74

sqrt_two = np.sqrt(2.)

class ParallelProgressBar:
    def __init__(self, title):
        self.title = title
        mylog.info(f"Starting '{title}'")

    def update(self, *args, **kwargs):
        return

    def close(self):
        mylog.info(f"Finishing '{self.title}'")

class SourceModel:
    def __init__(self, prng=None):
        self.spectral_norm = None
        self.redshift = None
        self.prng = parse_prng(prng)

    def __call__(self, chunk):
        pass

    def setup_model(self, data_source, redshift, spectral_norm):
        self.spectral_norm = spectral_norm
        self.redshift = redshift

    def cleanup_model(self):
        pass

metal_abund = {"angr": 0.0189,
               "aspl": 0.0134,
               "wilm": 0.0134,
               "lodd": 0.0133}

class ThermalSourceModel(SourceModel):
    r"""
    Initialize a source model from a thermal spectrum.

    Parameters
    ----------
    spectral_model : string
        The thermal model spectrum type to use. Known options are "apec".
    emin : float
        The minimum energy for the spectrum in keV.
    emax : float
        The maximum energy for the spectrum in keV.
    nchan : integer
        The number of channels in the spectrum.
    Zmet : float, string, or tuple of strings
        The metallicity. If a float, assumes a constant metallicity throughout
        in solar units. If a string or tuple of strings, is taken to be the
        name of the metallicity field.
    temperature_field : string or (ftype, fname) tuple, optional
        The yt temperature field to use for the thermal modeling. Must have
        units of Kelvin. If not specified, the default temperature field for
        the dataset will be used.
    emission_measure_field : string or (ftype, fname) tuple, optional
        The yt emission measure field to use for the thermal modeling. Must
        have units of cm^-3. If not specified, the default emission measure
        field for the dataset will be used or derived.
    kT_min : float, optional
        The default minimum temperature in keV to compute emission for.
        Default: 0.025
    kT_max : float, optional
        The default maximum temperature in keV to compute emission for.
        Default: 64.0
    n_kT : integer, optional
        The number of temperature bins to use when computing emission.
        Default: 10000
    kT_scale : string, optional
        The scaling of the bins to use when computing emission,
        "linear" or "log". Default: "linear"
    max_density : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The maximum density of the cells or particles to use when generating
        photons. If a float, the units are assumed to be g/cm**3.
        Default: 5e-25 g/cm**3.
    var_elem : dictionary, optional
        Elements that should be allowed to vary freely from the single abundance
        parameter. Each dictionary value, specified by the abundance symbol,
        corresponds to the abundance of that symbol. If a float, it is understood
        to be constant and in solar units. If a string or tuple of strings, it is
        assumed to be a spatially varying field. Default: None
    method : string, optional
        The method used to generate the photon energies from the spectrum:
        "invert_cdf": Invert the cumulative distribution function of the spectrum.
        "accept_reject": Acceptance-rejection method using the spectrum.
        The first method should be sufficient for most cases.
    thermal_broad : boolean, optional
        Whether or not the spectral lines should be thermally
        broadened. Default: True
    model_root : string, optional
        The directory root where the model files are stored. If not provided,
        a default location known to pyXSIM is used.
    model_vers : string, optional
        The version identifier string for the model files, e.g.
        "2.0.2". Default depends on the model used.
    nei : boolean, optional
        If True, use the non-equilibrium ionization tables. These are
        not supplied with pyXSIM/SOXS but must be downloaded separately, in
        which case the *apec_root* parameter must also be set to their
        location. Default: False
    nolines : boolean, optional
        Turn off lines entirely for generating emission.
        Default: False
    abund_table : string or array_like, optional
        The abundance table to be used for solar abundances.
        Either a string corresponding to a built-in table or an array
        of 30 floats corresponding to the abundances of each element
        relative to the abundance of H. Default is "angr".
        Built-in options are:
        "angr" : from Anders E. & Grevesse N. (1989, Geochimica et
        Cosmochimica Acta 53, 197)
        "aspl" : from Asplund M., Grevesse N., Sauval A.J. & Scott
        P. (2009, ARAA, 47, 481)
        "wilm" : from Wilms, Allen & McCray (2000, ApJ 542, 914
        except for elements not listed which are given zero abundance)
        "lodd" : from Lodders, K (2003, ApJ 591, 1220)
    prng : integer or :class:`~numpy.random.RandomState` object
        A pseudo-random number generator. Typically will only be specified
        if you have a reason to generate the same set of random numbers,
        such as for a test. Default is to use the :mod:`numpy.random` module.

    Examples
    --------
    >>> source_model = ThermalSourceModel("apec", 0.1, 10.0, 10000,
    ...                                   ("gas", "metallicity"))
    """
    def __init__(self, spectral_model, emin, emax, nchan, Zmet,
                 temperature_field=None, emission_measure_field=None,
                 kT_min=0.025, kT_max=64.0, n_kT=10000, kT_scale="linear",
                 max_density=5.0e-25, var_elem=None, method="invert_cdf",
                 thermal_broad=True, model_root=None, model_vers=None,
                 nei=False, nolines=False, abund_table="angr",
                 prng=None):
        if isinstance(spectral_model, str):
            if spectral_model not in thermal_models:
                raise KeyError(f"{spectral_model} is not a known thermal "
                               f"spectral model!")
            spectral_model = thermal_models[spectral_model]
        self.temperature_field = temperature_field
        self.Zmet = Zmet
        self.nei = nei
        self.ftype = None
        if var_elem is None:
            var_elem = {}
            var_elem_keys = None
            self.num_var_elem = 0
        else:
            var_elem_keys = list(var_elem.keys())
            self.num_var_elem = len(var_elem_keys)
        self.var_elem = var_elem
        self.spectral_model = spectral_model(emin, emax, nchan,
                                             var_elem=var_elem_keys,
                                             thermal_broad=thermal_broad,
                                             model_root=model_root,
                                             model_vers=model_vers,
                                             nolines=nolines, nei=nei,
                                             abund_table=abund_table)
        self.var_elem_keys = self.spectral_model.var_elem_names
        self.var_ion_keys = self.spectral_model.var_ion_names
        self.method = method
        self.prng = parse_prng(prng)
        self.kT_min = kT_min
        self.kT_max = kT_max
        self.kT_scale = kT_scale
        self.n_kT = n_kT
        self.spectral_norm = None
        self.redshift = None
        self.pbar = None
        self.kT_bins = None
        self.dkT = None
        self.emission_measure_field = emission_measure_field
        self.Zconvert = 1.0
        self.abund_table = abund_table
        self.atable = self.spectral_model.atable
        self.mconvert = {}
        if max_density is not None:
            if not isinstance(max_density, YTQuantity):
                if isinstance(max_density, tuple):
                    max_density = YTQuantity(max_density[0], max_density[1])
                else:
                    max_density = YTQuantity(max_density, "g/cm**3")
        self.max_density = max_density
        self.density_field = None  # Will be determined later
        self.tot_num_cells = 0  # Will be determined later
        self.ftype = "gas"

    def setup_model(self, data_source, redshift, spectral_norm):
        if self.emission_measure_field is None:
            self.emission_measure_field = \
                (self.ftype, 'emission_measure')
        try:
            ftype = data_source.ds._get_field_info(
                self.emission_measure_field).name[0]
        except YTFieldNotFound:
            raise RuntimeError(f"The {self.emission_measure_field} field is not "
                               "found. If you do not have species fields in "
                               "your dataset, you may need to set "
                               "default_species_fields='ionized' in the call "
                               "to yt.load().")
        self.ftype = ftype
        self.redshift = redshift
        if not self.nei and not isinstance(self.Zmet, float):
            Z_units = str(data_source.ds._get_field_info(self.Zmet).units)
            if Z_units in ["dimensionless", "", "code_metallicity"]:
                Zsum = (self.atable*atomic_weights)[metal_elem].sum()
                self.Zconvert = atomic_weights[1]/(Zsum*solar_H_abund)
            elif Z_units == "Zsun":
                self.Zconvert = 1.0
            else:
                raise RuntimeError(f"I don't understand metallicity "
                                   f"units of {Z_units}!")
        if self.num_var_elem > 0:
            for key, value in self.var_elem.items():
                if not isinstance(value, float):
                    if "^" in key:
                        elem = key.split("^")[0]
                    else:
                        elem = key
                    n_elem = elem_names.index(elem)
                    m_units = str(data_source.ds._get_field_info(value).units)
                    if m_units in ["dimensionless", "", "code_metallicity"]:
                        m = self.atable[n_elem]*atomic_weights[n_elem]
                        self.mconvert[key] = atomic_weights[1]/(m*solar_H_abund)
                    elif m_units == "Zsun":
                        self.mconvert[key] = 1.0
                    else:
                        raise RuntimeError(f"I don't understand units of "
                                           f"{m_units} for element {key}!")
        self.density_field = (ftype, "density")
        mylog.info(f"Using emission measure field "
                   f"'{self.emission_measure_field}'.")
        if self.temperature_field is None:
            self.temperature_field = (ftype, 'temperature')
        mylog.info(f"Using temperature field "
                   f"'{self.temperature_field}'.")
        self.spectral_model.prepare_spectrum(redshift, self.kT_min,
                                             self.kT_max)
        self.spectral_norm = spectral_norm
        if self.kT_scale == "linear":
            self.kT_bins = np.linspace(self.kT_min, self.kT_max,
                                       num=self.n_kT+1)
        elif self.kT_scale == "log":
            self.kT_bins = np.logspace(np.log10(self.kT_min),
                                       np.log10(self.kT_max),
                                       num=self.n_kT+1)
        self.dkT = np.diff(self.kT_bins)
        citer = data_source.chunks([], "io")
        num_cells = 0
        for chunk in parallel_objects(citer):
            num_cells += chunk[self.temperature_field].size
        self.tot_num_cells = comm.mpi_allreduce(num_cells)
        if parallel_capable:
            self.pbar = ParallelProgressBar("Processing cells/particles ")
        else:
            self.pbar = tqdm(leave=True, total=self.tot_num_cells,
                             desc="Processing cells/particles ")

    def cleanup_model(self):
        self.emission_measure_field = None
        self.temperature_field = None
        self.pbar.close()

    def __call__(self, chunk):

        num_photons_max = 10000000
        emid = self.spectral_model.emid
        ebins = self.spectral_model.ebins
        nchan = len(emid)

        orig_ncells = chunk[self.temperature_field].size
        if orig_ncells == 0:
            return
        if self.max_density is None:
            dens_cut = np.ones(orig_ncells, dtype="bool")
        else:
            dens_cut = chunk[self.density_field] < self.max_density
        kT = np.atleast_1d(
            chunk[self.temperature_field].to_value("keV", "thermal"))
        EM = np.atleast_1d(chunk[self.emission_measure_field].d*dens_cut)

        idxs = np.argsort(kT)

        kT_sorted = kT[idxs]
        idx_min = np.searchsorted(kT_sorted, self.kT_min)
        idx_max = np.searchsorted(kT_sorted, self.kT_max)
        idxs = idxs[idx_min:idx_max]
        num_cells = len(idxs)

        if num_cells == 0:
            self.pbar.update(orig_ncells)
            return
        else:
            self.pbar.update(orig_ncells-num_cells)

        kT_idxs = np.digitize(kT[idxs], self.kT_bins)-1
        bcounts = np.bincount(kT_idxs).astype("int")
        bcounts = bcounts[bcounts > 0]
        n = int(0)
        bcell = []
        ecell = []
        for bcount in bcounts:
            bcell.append(n)
            ecell.append(n+bcount)
            n += bcount
        kT_idxs = np.unique(kT_idxs)

        cell_em = EM[idxs]*self.spectral_norm

        if self.nei:
            metalZ = np.zeros(num_cells)
            elem_keys = self.var_ion_keys
        else:
            elem_keys = self.var_elem_keys
            if isinstance(self.Zmet, float):
                metalZ = self.Zmet*np.ones(num_cells)
            else:
                metalZ = np.atleast_1d(chunk[self.Zmet].d[idxs]*
                                       self.Zconvert)

        elemZ = None
        if self.num_var_elem > 0:
            elemZ = np.zeros((self.num_var_elem, num_cells))
            for j, key in enumerate(elem_keys):
                value = self.var_elem[key]
                if isinstance(value, float):
                    elemZ[j, :] = value
                else:
                    elemZ[j, :] = np.atleast_1d(chunk[value].d[idxs]*
                                                self.mconvert[key])

        number_of_photons = np.zeros(num_cells, dtype="int64")
        energies = np.zeros(num_photons_max)

        start_e = 0
        end_e = 0

        for ibegin, iend, bcount, ikT in zip(bcell, ecell, bcounts, kT_idxs):

            self.pbar.update(bcount)

            kT = self.kT_bins[ikT] + 0.5*self.dkT[ikT]

            cem = cell_em[ibegin:iend]

            cspec, mspec, vspec = self.spectral_model.get_spectrum(kT)

            tot_ph_c = cspec.d.sum()
            tot_ph_m = mspec.d.sum()

            cell_norm_c = tot_ph_c*cem
            cell_norm_m = tot_ph_m*metalZ[ibegin:iend]*cem
            cell_norm = cell_norm_c + cell_norm_m

            if vspec is not None:
                cell_norm_v = np.zeros(cem.size)
                for j in range(self.num_var_elem):
                    tot_ph_v = vspec.d[j, :].sum()
                    cell_norm_v += tot_ph_v*elemZ[j, ibegin:iend]*cem
                cell_norm += cell_norm_v

            cell_n = ensure_numpy_array(self.prng.poisson(lam=cell_norm))

            number_of_photons[ibegin:iend] = cell_n

            end_e += int(cell_n.sum())

            if self.method == "invert_cdf":
                cumspec_c = np.insert(np.cumsum(cspec.d), 0, 0.0)
                cumspec_m = np.insert(np.cumsum(mspec.d), 0, 0.0)
                if vspec is None:
                    cumspec_v = None
                else:
                    cumspec_v = np.zeros((self.num_var_elem, nchan+1))
                    for j in range(self.num_var_elem):
                        cumspec_v[j, 1:] = np.cumsum(vspec.d[j, :])

            ei = start_e
            for icell in range(ibegin, iend):
                cn = number_of_photons[icell]
                if cn == 0:
                    continue
                # The rather verbose form of the few next statements is a
                # result of code optimization and shouldn't be changed
                # without checking for perfomance degradation. See
                # https://bitbucket.org/yt_analysis/yt/pull-requests/1766
                # for details.
                if self.method == "invert_cdf":
                    cumspec = cumspec_c
                    cumspec += metalZ[icell] * cumspec_m
                    if cumspec_v is not None:
                        for j in range(self.num_var_elem):
                            cumspec += elemZ[j, icell]*cumspec_v[j, :]
                    norm_factor = 1.0 / cumspec[-1]
                    cumspec *= norm_factor
                    randvec = self.prng.uniform(size=cn)
                    randvec.sort()
                    cell_e = np.interp(randvec, cumspec, ebins)
                elif self.method == "accept_reject":
                    tot_spec = cspec.d
                    tot_spec += metalZ[icell] * mspec.d
                    if vspec is not None:
                        for j in range(self.num_var_elem):
                            tot_spec += elemZ[j, icell]*vspec.d[j, :]
                    norm_factor = 1.0 / tot_spec.sum()
                    tot_spec *= norm_factor
                    eidxs = self.prng.choice(nchan, size=cn, p=tot_spec)
                    cell_e = emid[eidxs]
                while ei+cn > num_photons_max:
                    num_photons_max *= 2
                if num_photons_max > energies.size:
                    energies.resize(num_photons_max, refcheck=False)
                energies[ei:ei+cn] = cell_e
                ei += cn

            start_e = end_e

        active_cells = number_of_photons > 0
        idxs = idxs[active_cells]
        ncells = idxs.size

        return ncells, number_of_photons[active_cells], idxs, energies[:end_e].copy()

    def cleanup_model(self):
        self.pbar.close()

class PowerLawSourceModel(SourceModel):
    r"""
    Initialize a source model from a power-law spectrum.

    Parameters
    ----------
    e0 : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The reference energy of the power law, in the rest frame of the source.
        If units are not given, they are assumed to be in keV.
    emin : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The minimum energy of the photons to be generated, in the rest frame of
        the source. If units are not given, they are assumed to be in keV.
    emax : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The maximum energy of the photons to be generated, in the rest frame of
        the source. If units are not given, they are assumed to be in keV.
    emission_field : string or (ftype, fname) tuple
        The field corresponding to the specific photon count rate per cell or
        particle, in the rest frame of the source, which serves as the
        normalization for the power law. Must be in counts/s/keV.
    index : float, string, or (ftype, fname) tuple
        The power-law index of the spectrum. Either a float for a single power law or
        the name of a field that corresponds to the power law.
    prng : integer or :class:`~numpy.random.RandomState` object
        A pseudo-random number generator. Typically will only be specified
        if you have a reason to generate the same set of random numbers, such as for a
        test. Default is to use the :mod:`numpy.random` module.

    Examples
    --------
    >>> e0 = (1.0, "keV")
    >>> emin = (0.01, "keV")
    >>> emax = (100., "keV")
    >>> plaw_model = PowerLawSourceModel(e0, emin, emax, ("gas", "norm"), ("gas", "index"))
    """
    def __init__(self, e0, emin, emax, emission_field, alpha, prng=None):
        self.e0 = parse_value(e0, "keV")
        self.emin = parse_value(emin, "keV")
        self.emax = parse_value(emax, "keV")
        self.emission_field = emission_field
        self.alpha = alpha
        self.prng = parse_prng(prng)
        self.spectral_norm = None
        self.redshift = None
        self.ftype = None

    def setup_model(self, data_source, redshift, spectral_norm):
        self.spectral_norm = spectral_norm
        self.redshift = redshift
        self.scale_factor = 1.0 / (1.0 + self.redshift)
        self.ftype = data_source.ds._get_field_info(self.emission_field).name[0]

    def __call__(self, chunk):

        num_cells = len(chunk[self.emission_field])

        if isinstance(self.alpha, float):
            alpha = self.alpha*np.ones(num_cells)
        else:
            alpha = chunk[self.alpha].v

        norm_fac = (self.emax.v**(1.-alpha)-self.emin.v**(1.-alpha))
        norm_fac[alpha == 1] = np.log(self.emax.v/self.emin.v)
        norm = norm_fac*chunk[self.emission_field].v*self.e0.v**alpha
        norm[alpha != 1] /= (1.-alpha[alpha != 1])
        norm *= self.spectral_norm*self.scale_factor

        number_of_photons = self.prng.poisson(lam=norm)

        energies = np.zeros(number_of_photons.sum())

        start_e = 0
        end_e = 0
        for i in range(num_cells):
            if number_of_photons[i] > 0:
                end_e = start_e+number_of_photons[i]
                u = self.prng.uniform(size=number_of_photons[i])
                if alpha[i] == 1:
                    e = self.emin.v*(self.emax.v/self.emin.v)**u
                else:
                    e = self.emin.v**(1.-alpha[i]) + u*norm_fac[i]
                    e **= 1./(1.-alpha[i])
                energies[start_e:end_e] = e * self.scale_factor
                start_e = end_e

        active_cells = number_of_photons > 0
        ncells = active_cells.sum()

        return ncells, number_of_photons[active_cells], active_cells, energies[:end_e].copy()

class LineSourceModel(SourceModel):
    r"""
    Initialize a source model from a single line.

    Parameters
    ----------
    e0 : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The location of the emission line in energy in the rest frame of the
        source. If units are not given, they are assumed to be in keV.
    emission_field : string or (ftype, fname) tuple
        The field corresponding to the photon count rate per cell or particle,
        in the rest frame of the source, which serves as the normalization for
        the line. Must be in counts/s.
    sigma : float, (value, unit) tuple, :class:`~yt.units.yt_array.YTQuantity`, or :class:`~astropy.units.Quantity`
        The standard intrinsic deviation of the emission line (not from Doppler
        broadening, which is handled in the projection step). Units of
        velocity or energy are accepted. If units are not given, they
        are assumed to be in keV. If set to a field name, the line broadening
        is assumed to be based on this field (in units of velocity or energy).
        If set to None (the default), it is assumed that the line is unbroadened.
    prng : integer or :class:`~numpy.random.RandomState` object
        A pseudo-random number generator. Typically will only be specified
        if you have a reason to generate the same set of random numbers, such as for a
        test. Default is to use the :mod:`numpy.random` module.

    Examples
    --------
    >>> location = (3.5, "keV")
    >>> sigma = (1000., "km/s")
    >>> line_model = LineEmissionSourceModel(location, "dark_matter_density_squared", sigma=sigma)
    """
    def __init__(self, e0, emission_field, sigma=None, prng=None):
        from unyt.exceptions import UnitConversionError
        self.e0 = parse_value(e0, "keV")
        if isinstance(sigma, Number):
            self.sigma = parse_value(sigma, "keV")
        elif isunitful(sigma):
            # The broadening is constant
            try:
                self.sigma = parse_value(sigma, "km/s")
                self.sigma *= self.e0/clight
                self.sigma.convert_to_units("keV")
            except UnitConversionError:
                self.sigma = parse_value(sigma, "keV")
        else:
            # Either no broadening or a field name
            self.sigma = sigma
        self.emission_field = emission_field
        self.prng = parse_prng(prng)
        self.spectral_norm = None
        self.redshift = None
        self.ftype = None

    def setup_model(self, data_source, redshift, spectral_norm):
        self.spectral_norm = spectral_norm
        self.redshift = redshift
        self.scale_factor = 1.0 / (1.0 + self.redshift)
        self.ftype = data_source.ds._get_field_info(self.emission_field).name[0]

    def __call__(self, chunk):
        num_cells = len(chunk[self.emission_field])
        F = chunk[self.emission_field]*self.spectral_norm*self.scale_factor
        number_of_photons = self.prng.poisson(lam=F.in_cgs().v)

        energies = self.e0*np.ones(number_of_photons.sum())

        if isinstance(self.sigma, YTQuantity):
            dE = self.prng.normal(loc=0.0, scale=float(self.sigma),
                                  size=number_of_photons.sum())*self.e0.uq
            energies += dE
        elif self.sigma is not None:
            sigma = (chunk[self.sigma]*self.e0/clight).in_units("keV")
            start_e = 0
            for i in range(num_cells):
                if number_of_photons[i] > 0:
                    end_e = start_e+number_of_photons[i]
                    dE = self.prng.normal(loc=0.0, scale=float(sigma[i]),
                                          size=number_of_photons[i])*self.e0.uq
                    energies[start_e:end_e] += dE
                    start_e = end_e

        energies = energies * self.scale_factor

        active_cells = number_of_photons > 0
        ncells = active_cells.sum()

        return ncells, number_of_photons[active_cells], active_cells, energies