# -*- coding: utf-8 -*-
"""
@author: NysanAskar
"""

import numpy as np
import tensorflow as tf
from keras import backend as K
import keras

from tensorflow.keras.layers import (
    Add,
    Input,
)
from utils import xywh_to_x1y1x2y2, broadcast_iou, binary_cross_entropy

anchors_wh = np.array([[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
                       [59, 119], [116, 90], [156, 198], [373, 326]],
                      np.float32) / 416

TOTAL_CLASSES = 80

backbone = [
    (32, 3, 1, 'same', 'mish', '0'),
    ["SPP", 1, '1'],
    (64, 3, 1, 'same', 'mish', '2'),
    ["CSP", 2, '3'],
    (128, 1, 1, 'same', 'mish', '4'),
    ["CSP", 8, '5'],
    (256, 1, 1, 'same', 'mish', '6'),
    ["CSP", 8, '7'],
    (512, 1, 1, 'same', 'mish', '8'),
    ["CSP", 4, '9'],
    (1024, 1, 1, 'same', 'mish', '10'),
]

neck = [
    (512, 1, 1, 'same', 'leaky', '0'),# input_3
    (1024, 3, 1, 'same', 'leaky', '1'),
    ["M", '2'],
    (512, 1, 1, 'same', 'leaky', '3'),
    (1024, 3, 1, 'same', 'leaky', '4'),
    (512, 1, 1, 'same', 'leaky', '5'),# output_3
    (256, 1, 1, 'same', 'leaky', '6'),
    "U", #'7'
    ["C",'8'],
    (256, 1, 1, 'same', 'mish', '9'),
    (512, 3, 1, 'same', 'mish', '10'),
    (256, 1, 1, 'same', 'mish', '11'),
    (512, 3, 1, 'same', 'mish', '12'),
    (256, 1, 1, 'same', 'mish', '13'),# output_2
    (128, 1, 1, 'same', 'leaky', '14'),
    "U", #'15'
    ["C",'16'],
    (128, 1, 1, 'same', 'leaky', '17'),
    (256, 3, 1, 'same', 'mish', '18'),
    (128, 1, 1, 'same', 'leaky', '19'),
    (256, 3, 1, 'same', 'mish', '20'),
    (128, 1, 1, 'same', 'leaky', '21'),
    (256, 3, 1, 'same', 'mish', '22'),# output_1
]

head = [
    ["S", 256, '0'],# input_3 -> output_1
    (256, 3, 2, 'valid', 'leaky', '1'),# input_3
    ["C", 256, '2'], # concatanate with input_2
    (512, 3, 1, 'same', 'leaky', '3'),
    (256, 1, 1, 'same', 'leaky', '4'),
    (512, 3, 1, 'same', 'leaky', '5'),
    (256, 1, 1, 'same', 'leaky', '6'),
    ["S", 256, '7'],# output_2
    (512, 3, 2, 'valid', 'leaky', '8'),
    ["C", 512, '9'], # concatanate with input_1
    (1024, 3, 1, 'same', 'leaky', '10'),
    (512, 1, 1, 'same', 'leaky', '11'),# output_2
    (1024, 3, 1, 'same', 'leaky', '12'),
    (512, 1, 1, 'same', 'leaky', '13'),
    ["S", 512, '14'],# output_3
]

class Mish(tf.keras.layers.Layer):
    '''
    Mish Activation Function.
    .. math::
        Mish(x) = x * tanh(softplus(x))
    '''
    def __init__(self, **kwargs):
        super(Mish, self).__init__(**kwargs)
    def call(self, inputs):
        return inputs * K.tanh(K.softplus(inputs))

class CNNBlock(tf.keras.layers.Layer):
    def __init__(self, out_channels, bn_act=True, padding='same', activation='leaky',**kwargs):
        super().__init__()
        self.padding = padding
        self.conv = tf.keras.layers.Conv2D(filters = out_channels, use_bias=not bn_act, padding=padding, **kwargs)
        self.bn = tf.keras.layers.BatchNormalization()
        self.leaky = tf.keras.layers.LeakyReLU(0.1)
        self.use_bn_act = bn_act
        self.zero_pad = tf.keras.layers.ZeroPadding2D(padding=(1, 1))
        self.activ_func = activation
        self.mish = Mish()
    def call(self, input_tensor):
        if self.activ_func == 'leaky':
            if self.padding == 'same':
                if self.use_bn_act:
                    return self.leaky(self.bn(self.conv(input_tensor)))
                else:
                    return self.conv(input_tensor)
            else:
                if self.use_bn_act:
                    z = self.zero_pad(input_tensor)
                    return self.leaky(self.bn(self.conv(z)))
                else:
                    z = self.zero_pad(input_tensor)
                    return self.conv(z) # TensorShape([3, 224, 224, 5])
        elif self.activ_func == 'mish':
            if self.padding == 'same':
                if self.use_bn_act:
                    return self.mish(self.bn(self.conv(input_tensor)))
                else:
                    return self.conv(input_tensor)
            else:
                if self.use_bn_act:
                    z = self.zero_pad(input_tensor)
                    return self.mish(self.bn(self.conv(z)))
                else:
                    z = self.zero_pad(input_tensor)
                    return self.conv(z) # TensorShape([3, 224, 224, 5])

class ResidualBlock(tf.keras.layers.Layer):
    def __init__(self, channels, use_residual=True, num_repeats=1):
        super().__init__()
        self.layers = []
        for _ in range(num_repeats):
            self.layers += [
                keras.Sequential([
                    CNNBlock(channels, kernel_size=1, activation='mish'),
                    CNNBlock(channels, kernel_size=3, padding='same', activation='mish')]
                )]
        self.use_residual = use_residual
        self.num_repeats = num_repeats
    def call(self, input_tensor):
        for layer in self.layers:
            if self.use_residual:
                x = Add()([input_tensor,layer(input_tensor)])
            else:
                x = layer(input_tensor)
        return x # TensorShape([3, 224, 224, 5])

class CSPBlock(tf.keras.layers.Layer):
    def __init__(self, channels, num_res_block=1):
        super().__init__()
        self.conv_1 = CNNBlock(out_channels=channels, padding='valid',
                             activation='mish', kernel_size=3, strides=2)
        self.conv_2 = CNNBlock(out_channels=channels//2, padding='same',
                             activation='mish', kernel_size=1, strides=1)
        self.conv_3 = CNNBlock(out_channels=channels//2, padding='same',
                             activation='mish', kernel_size=1, strides=1)
        self.res_block = ResidualBlock(channels=channels//2, num_repeats=num_res_block)
    def call(self, input_tensor):
        x = self.conv_1(input_tensor)
        x_1 = self.conv_2(x)
        x_2 = self.conv_2(x)
        x_2 = self.res_block(x_2)
        x_2 = self.conv_3(x_2)
        x = tf.concat([x_1, x_2], -1)
        return x # tf.Tensor: shape=(3, 208, 208, 256)

class backbone_layers(tf.keras.layers.Layer):
    def __init__(self, in_channels=3):
        super().__init__()
        self.in_channels = in_channels
        self.layers = self._create_conv_layers()
    def call(self, x):
        outputs = []  # for each scale
        for i, layer in enumerate(self.layers):
            if i in [6, 8, 10]:
                outputs.append(layer(x))
            x = layer(x)
        return outputs
    def _create_conv_layers(self):
        layers = []
        in_channels = self.in_channels
        for module in backbone:
            if isinstance(module, tuple):
                out_channels, kernel_size, strides, padding, ac_faunc, name_layer = module
                layers.append(
                    CNNBlock(
                        out_channels,
                        kernel_size=kernel_size,
                        strides=strides,
                        padding=padding, activation = ac_faunc, name = name_layer))
                in_channels = out_channels
            elif isinstance(module, list):
                num_repeats = module[1]
                layers.append(CSPBlock(in_channels*2, num_res_block=num_repeats))
        return layers

class max_pool(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.conv_1 = CNNBlock(out_channels=512, kernel_size=1,strides=1,
                        padding='same', activation = 'leaky')
        self.max_1 = tf.keras.layers.MaxPool2D((5, 5), strides=1, padding="same")
        self.max_2 = tf.keras.layers.MaxPool2D((9, 9), strides=1, padding="same")
        self.max_3 = tf.keras.layers.MaxPool2D((13, 13), strides=1, padding="same")
    def call(self, x):
        x_1 = self.conv_1(x)
        x_max_1 = self.max_1(x_1)
        x_max_2 = self.max_2(x_1)
        x_max_3 = self.max_3(x_1)
        x = tf.concat([x_1, x_max_1,x_max_2,x_max_3], -1)
        return x

class concat(tf.keras.layers.Layer):
    def __init__(self, channels, ):
        super().__init__()
        self.conc = tf.keras.layers.Concatenate(axis=-1)
        self.conv = CNNBlock(out_channels=channels, kernel_size=1,strides=1,
                        padding='same', activation = 'leaky')
    def call(self, x_1, x_2):
        x_2 = self.conv(x_2)
        x = self.conc([x_1, x_2,])
        return x

class yolov4_neck(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.layers = self._create_conv_layers()
    def call(self, input_tensor):
        outputs = []  # for each scale
        route_connections = []
        inputs = input_tensor[:2]
        x = input_tensor[2]
        for i, layer in enumerate(self.layers):
            if i in [5, 13, 22]:
                outputs.append(layer(x))
            elif isinstance(layer, tf.keras.layers.UpSampling2D):
                route_connections.append(layer(x))
                x = layer(x)
            elif isinstance(layer, concat):
                x = layer(route_connections.pop(), inputs.pop())
            else:
                x = layer(x)
        return outputs
    def _create_conv_layers(self):
        layers = []
        for module in neck:
            if isinstance(module, tuple):
                out_channels, kernel_size, strides, padding, act_func, name = module
                layers.append(
                    CNNBlock(
                        out_channels,
                        kernel_size=kernel_size,
                        strides=strides,
                        padding=padding, activation=act_func))
                in_channels = out_channels
            elif isinstance(module, list):
                l = module[0]
                if l == 'M':
                    layers.append(max_pool())
                elif l == 'C':
                    layers.append(concat(in_channels))
            elif isinstance(module, str):
                    layers.append(tf.keras.layers.UpSampling2D(size=2))
        return layers

class ScalePrediction(tf.keras.layers.Layer):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.conv_1 = CNNBlock(2 * in_channels, kernel_size=3, padding='same')
        self.conv_2 = CNNBlock((num_classes + 5) * 3, bn_act=False, kernel_size=1)
        self.num_classes = num_classes
    def call(self, input_tensor):
        if input_tensor.get_shape()[1] == 52:
            x = self.conv_2(input_tensor)
            x = tf.reshape(x, shape=(K.shape(x)[0], K.shape(x)[1], K.shape(x)[2], 3, self.num_classes + 5))
            return x
        else:
            x = self.conv_1(input_tensor)
            x = self.conv_2(x)
            x = tf.reshape(x, shape=(K.shape(x)[0], K.shape(x)[1], K.shape(x)[2], 3, self.num_classes + 5))
            return x # TensorShape([3, 416, 416, 3, 85])

class concat_head(tf.keras.layers.Layer):
    def __init__(self, in_channels):
        super().__init__()
        self.ch = in_channels
        self.conc = tf.keras.layers.Concatenate(axis=-1)
        self.conv = CNNBlock(self.ch, kernel_size=1,strides=1,
                        padding='same', activation = 'leaky')
    def call(self, x_1, x_2):
        x = self.conc([x_1, x_2,])
        x = self.conv(x)
        return x

class yolov4_head(tf.keras.layers.Layer):
    def __init__(self, num_classes):
        self.num_classes = num_classes
        super().__init__()
        self.layers = self._create_conv_layers()
    def call(self, input_tensor):
        outputs = []  # for each scale
        inputs = input_tensor[:2]
        x = input_tensor[2]
        for i, layer in enumerate(self.layers):
            if i in [0, 7, 14]:
                outputs.append(layer(x))
                continue
            elif isinstance(layer, concat_head):
                x = layer(x, inputs.pop())
            else:
                x = layer(x)
        return outputs
    def _create_conv_layers(self):
        layers = []
        for module in head:
            if isinstance(module, tuple):
                out_channels, kernel_size, strides, padding, act_func, name = module
                layers.append(
                    CNNBlock(
                        out_channels,
                        kernel_size=kernel_size,
                        strides=strides,
                        padding=padding, activation=act_func))
                in_channels = out_channels
            elif isinstance(module, list):
                l, c = module[0], module[1]
                if l == 'C':
                    layers.append(concat_head(in_channels=c))
                elif l == 'S':
                    layers.append(ScalePrediction(in_channels=c, num_classes= self.num_classes))
        return layers

def YoloV4(num_classes, shape=(416, 416, 3), training=True):
    inputs = Input(shape=shape)
    model_backbone = backbone_layers()
    out_backbone = model_backbone(inputs)
    model_neck = yolov4_neck()
    out_neck = model_neck(out_backbone)
    model_head = yolov4_head(num_classes)
    y_small, y_medium, y_large = model_head(out_neck)
    return tf.keras.Model(inputs, (y_small, y_medium, y_large))

def get_absolute_yolo_box(y_pred, valid_anchors_wh, num_classes):
    """
    inputs:
    y_pred: Prediction tensor from the model output, in the shape of (batch, grid, grid, anchor, 5 + num_classes)
    outputs:
    y_box: boxes in shape of (batch, grid, grid, anchor, 4), the last dimension is (xmin, ymin, xmax, ymax)
    objectness: probability that an object exists
    classes: probability of classes
    """
    t_xy, t_wh, objectness, classes = tf.split(
        y_pred, (2, 2, 1, num_classes), axis=-1)
    objectness = tf.sigmoid(objectness)
    classes = tf.sigmoid(classes)
    grid_size = tf.shape(y_pred)[1]
    C_xy = tf.meshgrid(tf.range(grid_size), tf.range(grid_size))
    C_xy = tf.stack(C_xy, axis=-1)
    C_xy = tf.expand_dims(C_xy, axis=2)  # [gx, gy, 1, 2]
    b_xy = tf.sigmoid(t_xy) + tf.cast(C_xy, tf.float32)
    b_xy = b_xy / tf.cast(grid_size, tf.float32)
    b_wh = tf.exp(t_wh) * valid_anchors_wh
    y_box = tf.concat([b_xy, b_wh], axis=-1)
    return y_box, objectness, classes

def get_relative_yolo_box(y_true, valid_anchors_wh):
    """
    This is the inverse of `get_absolute_yolo_box` above.
    """
    grid_size = tf.shape(y_true)[1]
    C_xy = tf.meshgrid(tf.range(grid_size), tf.range(grid_size))
    C_xy = tf.expand_dims(tf.stack(C_xy, axis=-1), axis=2)
    b_xy = y_true[..., 0:2]
    b_wh = y_true[..., 2:4]
    t_xy = b_xy * tf.cast(grid_size, tf.float32) - tf.cast(C_xy, tf.float32)
    t_wh = tf.math.log(b_wh / valid_anchors_wh)
    # b_wh could have some cells are 0, divided by anchor could result in inf or nan
    t_wh = tf.where(
        tf.logical_or(tf.math.is_inf(t_wh), tf.math.is_nan(t_wh)),
        tf.zeros_like(t_wh), t_wh)
    y_box = tf.concat([t_xy, t_wh], axis=-1)
    return y_box

class YoloLoss(object):
    def __init__(self, num_classes, valid_anchors_wh):
        self.num_classes = num_classes
        self.ignore_thresh = 0.5
        self.valid_anchors_wh = valid_anchors_wh
        self.lambda_coord = 5.0
        self.lamda_noobj = 0.5
    def __call__(self, y_true, y_pred):
        pred_xy_rel = tf.sigmoid(y_pred[..., 0:2]) # (None, grid, grid, 3,2)
        pred_wh_rel = y_pred[..., 2:4] # (None, grid, grid, 3,2)
        pred_box_abs, pred_obj, pred_class = get_absolute_yolo_box(
            y_pred, self.valid_anchors_wh, self.num_classes)
# # pred_box_abs (None, grid, grid, 3, 4), pred_obj (None, grid, grid, 3, 1), pred_class (None, grid, grid, 3, num_classes)
        pred_box_abs = xywh_to_x1y1x2y2(pred_box_abs)
        true_xy_abs, true_wh_abs, true_obj, true_class = tf.split(
            y_true, (2, 2, 1, self.num_classes), axis=-1)
# # true_xy_abs (None, grid, grid, 3, 2), true_wh_abs (None, grid, grid, 3, 2), true_obj (None, grid, grid, 3, 1), true_class (None, grid, grid, 3, num_classes)
        true_box_abs = tf.concat([true_xy_abs, true_wh_abs], axis=-1) #true_box_abs (None, grid, grid, 3, 4)
        true_box_abs = xywh_to_x1y1x2y2(true_box_abs)

        true_box_rel = get_relative_yolo_box(y_true, self.valid_anchors_wh)
        true_xy_rel = true_box_rel[..., 0:2]
        true_wh_rel = true_box_rel[..., 2:4]

        weight = 2 - true_wh_abs[..., 0] * true_wh_abs[..., 1] # (None, grid, grid, 3)

        xy_loss = self.calc_xy_loss(true_obj, true_xy_rel, pred_xy_rel, weight)
        wh_loss = self.calc_wh_loss(true_obj, true_wh_rel, pred_wh_rel, weight)
        class_loss, class_accuracy = self.calc_class_loss(true_obj, true_class, pred_class)

        ignore_mask, mIoU = self.calc_ignore_mask(true_obj, true_box_abs,
                                            pred_box_abs)
        obj_loss = self.calc_obj_loss(true_obj, pred_obj, ignore_mask)

        return xy_loss + wh_loss + class_loss + obj_loss, (xy_loss, wh_loss,
                                                           class_loss,
                                                           obj_loss,class_accuracy, mIoU)

    def calc_ignore_mask(self, true_obj, true_box, pred_box):
        # YOLOv3:
        true_box_shape = tf.shape(true_box)
        # (None, 13, 13, 3, 4)
        pred_box_shape = tf.shape(pred_box)
        # (None, 507, 4)
        true_box = tf.reshape(true_box, [true_box_shape[0], -1, 4])
        # sort true_box to have non-zero boxes rank first
        true_box = tf.sort(true_box, axis=1, direction="DESCENDING")
        # (None, 100, 4)
        true_box = true_box[:, 0:100, :]
        # (None, 507, 4)
        pred_box = tf.reshape(pred_box, [pred_box_shape[0], -1, 4])
        # (None, 507, 100)
        iou = broadcast_iou(pred_box, true_box)
        # (None, 507)
        best_iou = tf.reduce_max(iou, axis=-1)
        # (None, 13, 13, 3)
        best_iou = tf.reshape(best_iou, [pred_box_shape[0], pred_box_shape[1], pred_box_shape[2], pred_box_shape[3]])
        # ignore_mask = 1 => don't ignore
        # ignore_mask = 0 => should ignore
        ignore_mask = tf.cast(best_iou < self.ignore_thresh, tf.float32)
        # (None, 13, 13, 3, 1)
        ignore_mask = tf.expand_dims(ignore_mask, axis=-1)
        # (None, 13, 13, 3, 1)
        best_iou = tf.expand_dims(best_iou, axis=-1)
            # (None, )
        max_iou = tf.reduce_sum((best_iou*true_obj), axis=[1,2,3,4]) / tf.reduce_sum((true_obj[0])+0.001)
        return ignore_mask, max_iou

    def calc_obj_loss(self, true_obj, pred_obj, ignore_mask):
        """
        calculate loss of objectness: sum of L2 distances
        inputs:
        true_obj: objectness from ground truth in shape of (batch, grid, grid, anchor, num_classes)
        pred_obj: objectness from model prediction in shape of (batch, grid, grid, anchor, num_classes)
        outputs:
        obj_loss: objectness loss
        """
        obj_entropy = binary_cross_entropy(pred_obj, true_obj)

        obj_loss = true_obj * obj_entropy
        noobj_loss = (1 - true_obj) * obj_entropy * ignore_mask

        obj_loss = tf.reduce_sum(obj_loss, axis=(1, 2, 3, 4))
        noobj_loss = tf.reduce_sum(
            noobj_loss, axis=(1, 2, 3, 4)) * self.lamda_noobj
        return obj_loss + noobj_loss

    def calc_class_loss(self, true_obj, true_class, pred_class):
        """
        calculate loss of class prediction
        inputs:
        true_obj: if the object present from ground truth in shape of (batch, grid, grid, anchor, 1)
        true_class: one-hot class from ground truth in shape of (batch, grid, grid, anchor, num_classes)
        pred_class: one-hot class from model prediction in shape of (batch, grid, grid, anchor, num_classes)
        outputs:
        class_loss: class loss
        """
        pred_class_new = tf.cast(pred_class>0.5, tf.float32)
        class_accuracy = 1 - tf.reduce_sum(tf.abs(pred_class_new - true_class), axis=[1,2,3,4]) / tf.cast(tf.size(true_class[0]), tf.float32)
        class_loss = binary_cross_entropy(pred_class, true_class)
        class_loss = true_obj * class_loss
        class_loss = tf.reduce_sum(class_loss, axis=(1, 2, 3, 4))
        return class_loss, class_accuracy

    def calc_xy_loss(self, true_obj, true_xy, pred_xy, weight):
        """
        calculate loss of the centroid coordinate: sum of L2 distances
        inputs:
        true_obj: if the object present from ground truth in shape of (batch, grid, grid, anchor, 1)
        true_xy: centroid x and y from ground truth in shape of (batch, grid, grid, anchor, 2)
        pred_xy: centroid x and y from model prediction in shape of (batch, grid, grid, anchor, 2)
        weight: weight adjustment, reward smaller bounding box
        outputs:
        xy_loss: centroid loss
        """
        # shape (batch, grid, grid, anchor), eg. (32, 13, 13, 3)
        xy_loss = tf.reduce_sum(tf.square(true_xy - pred_xy), axis=-1)

        # in order to element-wise multiply the result from tf.reduce_sum
        # we need to squeeze one dimension for objectness here
        true_obj = tf.squeeze(true_obj, axis=-1)
        xy_loss = true_obj * xy_loss * weight
        xy_loss = tf.reduce_sum(xy_loss, axis=(1, 2, 3)) * self.lambda_coord
        return xy_loss # (batch_size/None, )

    def calc_wh_loss(self, true_obj, true_wh, pred_wh, weight):
        """
        calculate loss of the width and height: sum of L2 distances
        inputs:
        true_obj: if the object present from ground truth in shape of (batch, grid, grid, anchor, 1)
        true_wh: width and height from ground truth in shape of (batch, grid, grid, anchor, 2)
        pred_wh: width and height from model prediction in shape of (batch, grid, grid, anchor, 2)
        weight: weight adjustment, reward smaller bounding box
        outputs:
        wh_loss: width and height loss
        """
        # shape (batch, grid, grid, anchor), eg. (32, 13, 13, 3)
        wh_loss = tf.reduce_sum(tf.square(true_wh - pred_wh), axis=-1)
        true_obj = tf.squeeze(true_obj, axis=-1)
        wh_loss = true_obj * wh_loss * weight
        wh_loss = tf.reduce_sum(wh_loss, axis=(1, 2, 3)) * self.lambda_coord
        return wh_loss # (batch_size/None,)


"""
module for generating C, C++, Fortran77, Fortran90 and python routines that
evaluate sympy expressions. This module is work in progress. Only the
milestones with a '+' character in the list below have been completed.

--- How is sympy.utilities.codegen different from sympy.printing.ccode? ---

We considered the idea to extend the printing routines for sympy functions in
such a way that it prints complete compilable code, but this leads to a few
unsurmountable issues that can only be tackled with dedicated code generator:

- For C, one needs both a code and a header file, while the printing routines
generate just one string. This code generator can be extended to support .pyf
files for f2py.

- Sympy functions are not concerned with programming-technical issues, such as
input, output and input-output arguments. Other examples are contiguous or
non-contiguous arrays, including headers of other libraries such as gsl or others.

- It is highly interesting to evaluate several sympy functions in one C routine,
eventually sharing common intermediate results with the help of the cse routine.
This is more than just printing.

- From the programming perspective, expressions with constants should be
evaluated in the code generator as much as possible. This is different for
printing.

--- Basic assumptions ---

* A generic Routine data structure describes the routine that must be translated
  into C/Fortran/... code. This data structure covers all features present in
  one or more of the supported languages.

* Descendants from the CodeGen class transform multiple Routine instances into
  compilable code. Each derived class translates into a specific language.

* In many cases, one wants a simple workflow. The friendly functions in the last
  part are a simple api on top of the Routine/CodeGen stuff. They are easier to
  use, but are less powerful.

--- Milestones ---

+ First working version with scalar input arguments, generating C code, tests
+ Friendly functions that are easier to use than the rigorous Routine/CodeGen
  workflow.
+ Integer and Real numbers as input and output
- Optional extra include lines for libraries/objects that can eval special
  functions
- Test other C compilers and libraries: gcc, tcc, libtcc, gcc+gsl, ...
- Output arguments
- InputOutput arguments
- Sort input/output arguments properly
- Contiguous array arguments (sympy matrices)
- Non-contiguous array arguments (sympy matrices)
- ccode must raise an error when it encounters something that can not be
  translated into c. ccode(integrate(sin(x)/x, x)) does not make sense.
- Complex numbers as input and output
- Also generate .pyf code for f2py
- A default complex datatype
- Include extra information in the header: date, user, hostname, sha1 hash, ...
- Isolate constants and evaluate them beforehand in double precision
- Common Subexpression Elimination
- User defined comments in the generated code
- Fortran 77
- Fortran 90
- C++
- Python
- ...
"""

from sympy.core.symbol import Symbol
from sympy.core.expr import Expr
from sympy.core.containers import Tuple
from sympy.printing.ccode import ccode
from sympy.printing.fcode import fcode
from sympy.tensor import Idx, IndexedElement
from sympy.core.relational import Equality
from sympy.utilities import flatten

from StringIO import StringIO
import sympy
import os

__all__ = [
    # description of routines
    "Routine", "DataType", "default_datatypes", "get_default_datatype",
    "Argument", "InputArgument", "Result",
    # routines -> code
    "CodeGen", "CCodeGen", "FCodeGen",
    # friendly functions
    "codegen",
]

#
# Description of routines
#

class Routine(object):
    """Generic description of an evaluation routine for a set of sympy expressions.

       A CodeGen class can translate instances of this class into C/Fortran/...
       code. The routine specification covers all the features present in these
       languages. The CodeGen part must raise an exception when certain features
       are not present in the target language. For example, multiple return
       values are possible in Python, but not in C or Fortran. Another example:
       Fortran and Python support complex numbers, while C does not.
    """
    def __init__(self, name, expr):
        """Initialize a Routine instance.

           Arguments:
             name  --  A string with the name of this routine in the generated
                       code
             expr  --  The sympy expression that the Routine instance will represent.
                       If given a list or tuple of expressions, the routine
                       will be considered to have multiple return values.

            A decision about whether to use output arguments or return values,
            is made depending on the mathematical expressions.  For an expression
            of type Equality, the left hand side is made into an OutputArgument
            (or an InOutArgument if appropriate).  Else, the calculated
            expression is the return values of the routine.

            A tuple of exressions can be used to create a routine with both
            return value(s) and output argument(s).

        """
        arg_list = []

        if isinstance(expr, (list, tuple)):
            if not expr:
                raise ValueError("No expression given")
            expressions = Tuple(*expr)
        else:
            expressions = Tuple(expr)

        # local variables
        local_vars = set([i.label for i in expressions.atoms(Idx)])

        # symbols that should be arguments
        symbols = expressions.atoms(Symbol) - local_vars

        # Decide whether to use output argument or return value
        return_val = []
        output_args = []
        for expr in expressions:
            if isinstance(expr, Equality):
                out_arg = expr.lhs
                expr = expr.rhs
                if isinstance(out_arg, IndexedElement):
                    dims = out_arg.dimensions
                    symbol = out_arg.stem.label
                elif isinstance(out_arg, Symbol):
                    dims = []
                    symbol = out_arg
                else:
                    raise CodeGenError("Only IndexedElement or Symbol can define output arguments")

                if expr.has(symbol):
                    output_args.append(InOutArgument(symbol, out_arg, expr, dimensions=dims))
                else:
                    output_args.append(OutputArgument(symbol, out_arg, expr, dimensions=dims))

                # avoid duplicate arguments
                symbols.remove(symbol)
            else:
                return_val.append(Result(expr))

        # setup input argument list
        array_symbols = {}
        for array in expressions.atoms(IndexedElement):
            array_symbols[array.stem.label] = array

        for symbol in sorted(symbols):
            if symbol in array_symbols:
                dims = []
                array = array_symbols[symbol]
                for i in array.indices:
                    if i.lower == None:
                        dims.append((S.Zero, S.Zero))
                    else:
                        dims.append((i.lower, i.upper))
                metadata = {'dimensions': dims}
            else:
                metadata = {}

            arg_list.append(InputArgument(symbol, **metadata))

        arg_list.extend(output_args)

        self.name = name
        self.arguments = arg_list
        self.results = return_val
        self.local_vars = local_vars

    @property
    def result_variables(self):
        """Returns a tuple of OutputArgument, InOutArgument and Result."""
        args = [arg for arg in self.arguments if isinstance(arg, (OutputArgument, InOutArgument))]
        args.extend(self.results)
        return args

class DataType(object):
    """Holds strings for a certain datatype in different programming languages."""
    def __init__(self, cname, fname, pyname):
        self.cname = cname
        self.fname = fname
        self.pyname = pyname

default_datatypes = {
    "int": DataType("int", "INTEGER*4", "int"),
    "float": DataType("double", "REAL*8", "float")
}

def get_default_datatype(expr):
    """Derives a decent data type based on the assumptions on the expression."""
    if expr.is_integer:
        return default_datatypes["int"]
    else:
        return default_datatypes["float"]

class Argument(object):
    """An abstract Argument data structure: a name and a data type.

       This structure is refined in the descendants below.
    """

    def __init__(self, name, datatype=None, dimensions=None, precision=None):
        """Initialize an input argument.

           name  --  must be of class Symbol
           datatype  --  When not given, the data type will be guessed based
                         on the assumptions on the symbol argument.
           dimension  --  If present, the argument is interpreted as an array.
                          Dimensions must be a sequence containing tuples
                          of first and last index in the range.
           precision  --  FIXME
        """

        if not isinstance(name, Symbol):
            raise TypeError("The first argument must be a sympy symbol.")
        if datatype is None:
            datatype = get_default_datatype(name)
        elif not isinstance(datatype, DataType):
            raise TypeError("The (optional) `datatype' argument must be an instance of the DataType class.")
        if dimensions and not isinstance(dimensions, (tuple, list)):
            raise TypeError("The dimension argument must be a sequence of tuples")
        self.name = name
        self.datatype = datatype
        self.dimensions = dimensions
        self.precision = precision

    def get_symbols(self):
        """Returns a set of all symbols related to this argument.

        Scalar arguments return themselves in a set, while array arguments return
        the array variable as well as all symbols the specifiy dimensions.
        """
        if self.dimensions:
            symbs = set(flatten(self.dimensions))
            symbs.add(self.name)
            return symbs
        else:
            return set([self.name])

class InputArgument(Argument):
    pass

class ResultBase(object):

    @property
    def needs_initialization(self):
        return self._need_initialization

    def _prepare_expr(self):
        """Depending on the expression, it may need to be prepared for loops

        Example:

        The math expression denoting a matrix vector product is

        y(i) = A(i,j)*x(j)          (imlicit summation over j)

        Obviously, the code that calculates the matrix vector product must change
        the right hand side a little in order to store intermediate results
        correctly ....................
                                      |
        y = 0                         |
        do i=1,m                      |
            do j = 1,n                V
                y(i) = A(i,j)*x(i) + y(i)
            end do j
        end do i

        This would not be necessary if the left hand side has all indices that are on
        the right hand side.  E.g. for a dyadic product:

        do i=1,m
            do j = 1,n
                A(i,j) = x(i)*y(j)
            end do j
        end do i

        """
        rhs_loops = set([ i.label for i in self.expr.atoms(Idx) ])
        lhs_loops = set([ i.label for i in self.result_var.atoms(Idx) ])
        if rhs_loops - lhs_loops:
            self.expr = self.expr + self.result_var

class OutputArgument(Argument, ResultBase):
    """OutputArgument are always initialized in the routine
    """
    _need_initialization = True
    def __init__(self, name, result_var, expr, datatype=None, dimensions=None, precision=None):
        Argument.__init__(self, name, datatype, dimensions, precision)
        self.expr = expr
        self.result_var = result_var
        self._prepare_expr()

class InOutArgument(Argument, ResultBase):
    """InOutArgument are never initialized in the routine
    """
    _need_initialization = False

    def __init__(self, name, result_var, expr, datatype=None, dimensions=None, precision=None):
        Argument.__init__(self, name, datatype, dimensions, precision)
        self.expr = expr
        self.result_var = result_var

class Result(ResultBase):
    """An expression for a scalar return value.

       The name result is used to avoid conflicts with the reserved word
       'return' in the python language. It is also shorter than ReturnValue.
    """
    _need_initialization = False

    def __init__(self, expr, datatype=None):
        """Initialize a (scalar) return value.

           The second argument is optional. When not given, the data type will
           be guessed based on the assumptions on the expression argument.
        """
        if not isinstance(expr, Expr):
            raise TypeError("The first argument must be a sympy expression.")
        if datatype is None:
            datatype = get_default_datatype(expr)
        elif not isinstance(datatype, DataType):
            raise TypeError("The (optional) second argument must be an instance of the DataType class.")
        self.expr = expr
        self.datatype = datatype
        self.result_var = Symbol('result_%s'%hash(expr))
        self._prepare_expr()

#
# Transformation of routine objects into code
#

class CodeGen(object):
    """Abstract class for the code generators."""

    def __init__(self, project="project"):
        """Initialize a code generator.

           Derived classes will offer more options that affect the generated
           code.
        """
        self.project = project

    def write(self, routines, prefix, to_files=False, header=True, empty=True):
        """Writes all the source code files for the given routines.

           The generate source is returned as a list of (filename, contents)
           tuples, or is written to files (see options). Each filename consists
           of the given prefix, appended with an appropriate extension.

           Arguments:
             routines  --  A list of Routine instances to be written
             prefix  --  The prefix for the output files

           Optional arguments:
             to_files  --  When True, the output is effectively written to
                           files. [DEFAULT=False] Otherwise, a list of
                           (filename, contents) tuples is returned.
             header  --  When True, a header comment is included on top of each
                         source file. [DEFAULT=True]
             empty  --  When True, empty lines are included to structure the
                        source files. [DEFAULT=True]
        """
        if to_files:
            for dump_fn in self.dump_fns:
                filename = "%s.%s" % (prefix, dump_fn.extension)
                f = file(filename, "w")
                dump_fn(self, routines, f, prefix, header, empty)
                f.close()
        else:
            result = []
            for dump_fn in self.dump_fns:
                filename = "%s.%s" % (prefix, dump_fn.extension)
                contents = StringIO()
                dump_fn(self, routines, contents, prefix, header, empty)
                result.append((filename, contents.getvalue()))
            return result

    def dump_code(self, routines, f, prefix, header=True, empty=True):
        """Write the code file by calling language specific methods in correct order

           The generated file contains all the definitions of the routines in
           low-level code and refers to the header file if appropriate.

           Arguments:
             routines  --  a list of Routine instances
             f  --  a file-like object to write the file to
             prefix  --  the filename prefix, used to refer to the proper header
                         file. Only the basename of the prefix is used.

           Optional arguments:
             header  --  When True, a header comment is included on top of each
                         source file. [DEFAULT=True]
             empty  --  When True, empty lines are included to structure the
                        source files. [DEFAULT=True]
        """

        code_lines = []
        code_lines.extend(self._preprosessor_statements(prefix))
        for routine in routines:
            if empty: code_lines.append("\n")
            code_lines.extend(self._get_routine_opening(routine))
            code_lines.extend(self._declare_arguments(routine))
            code_lines.extend(self._declare_locals(routine))
            if empty: code_lines.append("\n")
            code_lines.extend(self._call_printer(routine))
            if empty: code_lines.append("\n")
            code_lines.extend(self._get_routine_ending(routine))

        if header and code_lines:
            self._dump_header(f)
        print >> f, ''.join(code_lines),

class CodeGenError(Exception):
    pass

header_comment = """Code generated with sympy %(version)s

See http://www.sympy.org/ for more information.

This file is part of '%(project)s'
"""

class CCodeGen(CodeGen):

    def _dump_header(self, f):
        """Writes a common header for the generated files."""
        print >> f, "/" + "*"*78
        tmp = header_comment % {"version": sympy.__version__, "project": self.project}
        for line in tmp.splitlines():
            print >> f, " *%s* " % line.center(76)
        print >> f, " " + "*"*78 + "/"

    def get_prototype_result(self, routine):
        """Returns a string for the function prototype for the given routine and
           a single result object, which can be None.

           If the routine has multiple result objects, an CodeGenError is
           raised.

           See: http://en.wikipedia.org/wiki/Function_prototype
        """
        prototype = []
        if len(routine.results) > 1:
            raise CodeGenError("C only supports a single or no return value.")
        elif len(routine.results) == 1:
            result = routine.results[0]
            prototype.append(result.datatype.cname)
        else:
            result = None
            prototype.append("void")
        # name of the routine + arguments + curly opening brackets
        prototype.append("%s(%s)" % (
            routine.name,
            ", ".join("%s %s" % (arg.datatype.cname, arg.name) for arg in routine.arguments)
        ))
        return " ".join(prototype), result

    def _preprosessor_statements(self, prefix):
        code_lines = []
        code_lines.append("#include \"%s.h\"\n" % os.path.basename(prefix))
        code_lines.append("#include <math.h>\n")
        return code_lines

    def _get_routine_opening(self, routine):
        prototype, result = self.get_prototype_result(routine)
        return ["%s {\n" % prototype]

    def _declare_arguments(self, routine):
        # arguments are in prototype
        return []

    def _declare_locals(self, routine):
        code_list = []
        for var in sorted(routine.local_vars, key=str):
            typeinfo = get_default_datatype(var)
            code_list.append("%s %s\n" % (typeinfo.cname, var))
        return code_list

    def _call_printer(self, routine):
        prototype, result = self.get_prototype_result(routine)
        if result is not None:
            return ["  return %s;\n" % ccode(result.expr)]
        else:
            return []

    def _get_routine_ending(self, routine):
        return ["}\n"]

    def dump_c(self, routines, f, prefix, header=True, empty=True):
        self.dump_code(routines, f, prefix, header, empty)
    dump_c.extension = "c"

    def dump_h(self, routines, f, prefix, header=True, empty=True):
        """Writes the C header file.

           This file contains all the function declarations.

           Arguments:
             routines  --  a list of Routine instances
             f  --  a file-like object to write the file to
             prefix  --  the filename prefix, used to construct the include
                         guards.

           Optional arguments:
             header  --  When True, a header comment is included on top of each
                         source file. [DEFAULT=True]
             empty  --  When True, empty lines are included to structure the
                        source files. [DEFAULT=True]
        """
        if header:
            self._dump_header(f)
        guard_name = "%s__%s__H" % (self.project.replace(" ", "_").upper(), prefix.replace("/", "_").upper())
        # include guards
        if empty: print >> f
        print >> f, "#ifndef %s" % guard_name
        print >> f, "#define %s" % guard_name
        if empty: print >> f
        # declaration of the function prototypes
        for routine in routines:
            prototype, result = self.get_prototype_result(routine)
            print >> f, "%s;" % prototype
        # end if include guards
        if empty: print >> f
        print >> f, "#endif"
        if empty: print >> f
    dump_h.extension = "h"

    # This list of dump functions is used by CodeGen.write to know which dump
    # functions it has to call.
    dump_fns = [dump_c, dump_h]

class FCodeGen(CodeGen):
    """
    Generator for Fortran 95 code
    """

    def __init__(self, project='project'):
        CodeGen.__init__(self, project)

    def _dump_header(self, f):
        """Writes a common header for the generated files."""
        print >> f, "!" + "*"*78
        tmp = header_comment % {"version": sympy.__version__, "project": self.project}
        for line in tmp.splitlines():
            print >> f, "!*%s* " % line.center(76)
        print >> f, "!" + "*"*78

    def _preprosessor_statements(self, prefix):
        return []

    def _get_routine_opening(self, routine):
        """
        Returns the opening statements of the fortran routine
        """
        code_list = []
        if len(routine.results) > 1:
            raise CodeGenError("Fortran only supports a single or no return value.")
        elif len(routine.results) == 1:
            result = routine.results[0]
            code_list.append(result.datatype.fname)
            code_list.append("function")
        else:
            code_list.append("subroutine")

        # name of the routine + arguments
        code_list.append("%s(%s)\n" % (routine.name,
            ", ".join("%s" % arg.name for arg in routine.arguments)))
        code_list = [ " ".join(code_list) ]

        code_list.append('implicit none\n')
        return code_list

    def _declare_arguments(self, routine):
        # argument type declarations
        code_list = []
        array_list = []
        scalar_list = []
        for arg in routine.arguments:

            if isinstance(arg, InputArgument):
                typeinfo = "%s, intent(in)" % arg.datatype.fname
            elif isinstance(arg, InOutArgument):
                typeinfo = "%s, intent(inout)" % arg.datatype.fname
            elif isinstance(arg, OutputArgument):
                typeinfo = "%s, intent(out)" % arg.datatype.fname
            else:
                raise CodeGenError("Unkown Argument type: %s"%type(arg))

            if arg.dimensions:
                # fortran arrays start at 1
                dimstr = ", ".join(["%s:%s"%(dim[0]+1, dim[1]+1)
                    for dim in arg.dimensions])
                typeinfo += ", dimension(%s)" % dimstr
                array_list.append("%s :: %s\n" % (typeinfo, arg.name))
            else:
                scalar_list.append("%s :: %s\n" % (typeinfo, arg.name))

        # scalars first, because they can be used in array declarations
        code_list.extend(scalar_list)
        code_list.extend(array_list)

        return code_list

    def _declare_locals(self, routine):
        code_list = []
        for var in sorted(routine.local_vars, key=str):
            typeinfo = get_default_datatype(var)
            code_list.append("%s :: %s\n" % (typeinfo.fname, var))

        return code_list

    def _get_routine_ending(self, routine):
        """
        Returns the closing statements of the fortran routine
        """
        if len(routine.results) == 1:
            return ["end function\n"]
        else:
            return ["end subroutine\n"]

    def get_interface(self, routine):
        """Returns a string for the function interface for the given routine and
           a single result object, which can be None.

           If the routine has multiple result objects, a CodeGenError is
           raised.

           See: http://en.wikipedia.org/wiki/Function_prototype

        """
        prototype = [ "interface\n" ]
        prototype.extend(self._get_routine_opening(routine))
        prototype.extend(self._declare_arguments(routine))
        prototype.extend(self._get_routine_ending(routine))
        prototype.append("end interface\n")

        return "".join(prototype)

    def _get_result(self, routine):
        """Returns a single result object, which can be return value or outargument
        """

        if len(routine.results) > 1:
            raise CodeGenError("Fortran only supports a single or no return value.")
        elif len(routine.results) == 1:
            result = routine.results[0]
        else:
            outargs = [ arg for arg in routine.arguments if
                    isinstance(arg, (OutputArgument, InOutArgument))]
            if len(outargs) == 1:
                result = outargs[0]
            else:
                raise CodeGenError("FIXME: need one and only one result object. Got %s"%len(outargs))

        return result

    def _init_resultvars(self, routine):
        """Returns codelines that intialize the result variables if applicable.
        """
        code_lines = []
        for arg in routine.arguments:
            if isinstance(arg, OutputArgument):
                if arg.datatype.fname == 'REAL*8':
                    code_lines.append("%s = 0.d0\n" % arg.name)
                elif arg.datatype.fname == 'INTEGER*4':
                    code_lines.append("%s = 0\n" % arg.name)
                else:
                    raise NotImplementedError
        if routine.results:
            code_lines.append("%s = 0.d0\n" % routine.name)

        return code_lines

    def _call_printer(self, routine):
        code_lines = []
        for result in routine.result_variables:
            if isinstance(result, Result):
                assign_to = routine.name
            elif isinstance(result, (OutputArgument, InOutArgument)):
                assign_to = result.result_var

            constants, not_fortran, f_expr = fcode(result.expr,
                assign_to=assign_to, source_format='free', human=False)
            code_lines.extend(constants)
            if result.needs_initialization:
                code_lines.extend(self._init_resultvars(routine))
            code_lines.append("%s\n" % f_expr)
        return code_lines

    def dump_f95(self, routines, f, prefix, header=True, empty=True):
        self.dump_code(routines, f, prefix, header, empty)
    dump_f95.extension = "f90"

    def dump_h(self, routines, f, prefix, header=True, empty=True):
        """Writes the interface  header file.

           This file contains all the function declarations.

           Arguments:
             routines  --  a list of Routine instances
             f  --  a file-like object to write the file to
             prefix  --  the filename prefix, used to construct the include
                         guards.

           Optional arguments:
             header  --  When True, a header comment is included on top of each
                         source file. [DEFAULT=True]
             empty  --  When True, empty lines are included to structure the
                        source files. [DEFAULT=True]
        """
        if header:
            self._dump_header(f)
        if empty: print >> f
        # declaration of the function prototypes
        for routine in routines:
            prototype  = self.get_interface(routine)
            print >> f, prototype,
        if empty: print >> f
    dump_h.extension = "h"

    # This list of dump functions is used by CodeGen.write to know which dump
    # functions it has to call.
    dump_fns = [dump_f95, dump_h]

def get_code_generator(language, project):
    CodeGenClass = {"C": CCodeGen, "F95": FCodeGen}.get(language.upper())
    if CodeGenClass is None:
        raise ValueError("Language '%s' is not supported." % language)
    return CodeGenClass(project)

#
# Friendly functions
#

def codegen(name_expr, language, prefix, project="project", to_files=False, header=True, empty=True):
    """Write source code for the given expressions in the given language.

       Mandatory Arguments:
         name_expr  --  A single (name, expression) tuple or a list of
                        (name, expression) tuples. Each tuple corresponds to a
                        routine.  If the expression is an equality (an
                        insance of class Equality) the left hand side is considered
                        an output argument.
         language  --  A string that indicates the source code language. This
                       is case insensitive. For the moment, only 'C' is
                       supported.
         prefix  --  A prefix for the names of the files that contain the source
                     code. Proper (language dependent) suffixes will be
                     appended.

       Optional Arguments:
         project  --  A project name, used for making unique preprocessor
                      instructions. [DEFAULT="project"]
         to_files  --  When True, the code will be written to one or more files
                       with the given prefix, otherwise strings with the names
                       and contents of these files are returned. [DEFAULT=False]
         header  --  When True, a header is written on top of each source file.
                     [DEFAULT=True]
         empty  --  When True, empty lines are used to structure the code.
                    [DEFAULT=True]

       >>> from sympy import symbols
       >>> from sympy.utilities.codegen import codegen
       >>> from sympy.abc import x, y, z
       >>> [(c_name, c_code), (h_name, c_header)] = \\
       ...     codegen(("f", x+y*z), "C", "test", header=False, empty=False)
       >>> print c_name
       test.c
       >>> print c_code,
       #include "test.h"
       #include <math.h>
       double f(double x, double y, double z) {
         return x + y*z;
       }
       >>> print h_name
       test.h
       >>> print c_header,
       #ifndef PROJECT__TEST__H
       #define PROJECT__TEST__H
       double f(double x, double y, double z);
       #endif

    """

    # Initialize the code generator.
    code_gen = get_code_generator(language, project)

    # Construct the routines based on the name_expression pairs.
    #  mainly the input arguments require some work
    routines = []
    if isinstance(name_expr[0], basestring):
        # single tuple is given, turn it into a singleton list with a tuple.
        name_expr = [name_expr]

    for name, expr in name_expr:
        routines.append(Routine(name, expr))

    # Write the code.
    return code_gen.write(routines, prefix, to_files, header, empty)


#!/usr/bin/env python
# coding: utf-8

# # Lab 3: Building a Photometric Pipeline
#
# In this lab, we'll be using classes and functions to build a pipeline which will automatically extract the fluxes of stars in an image. We're all familiar with aperture photometry, but in this case, we're going to take the additional step of convolving our image with a PSF.
#
# Our pipeline will be split into several steps.
#
# 1. Reading in an image
# 2. Finding the local peaks in the image (the stars)
# 3. Calculating the centroid of each peak region
# 4. Convolving with the PSF and extracting flux.
#
#

# ```{warning}
# When you modify a class and re-run it, you need to re-instantiate (or make a new instance) of the class for it to reflect the changes. This is a common mistake in notebooks, where a class might get changed but your instance of it didn't get re-created. Watch out for this during the lab -- the safest bet is to just copy-paste your instantiation code (`a = MyObj()` type thing) into each problem so you sure you're working with the latest version of your class.
#
# You'll be copying and pasting (and overwriting) your class in this assignment quite a bit. That's not super realistic to reality (where you'd pick one cell and keep modifying it there), but for grading purposes we ask you follow along so we can give credit for each method of the class you complete and demonstrate along the way.
# ```

# ## Assignment Rubric
# ````{panels}
# The points for this assignment are as follows:
# - Problem 1 (4 pts)
# - Problem 2 (8 pts)
# - Problem 3 (5 pts)
# - Problem 4 (3 pts)
# ````

# ## Problem 1
# In this problem, we'll load in the image(s) for use and begin constructing our `PSFPhot` class.
#
# ### Problem 1.1
# Create a script (`.py` file) in this directory and copy into it your `load_fits()` and final `implot()` functions from last lab. Then import them here.
#
# Use your `load_fits()` to read in in the image `2020-04-15-0001.fits` and plot it with `implot()`.
#
# This image was taken of the M81/M82 field using the Dragonfly Telephoto Array (in a narrowband configuration).

# In[ ]:

# import your functions here

# In[1]:

# Use them to load and show the image we'll be working with

# ```{hint}
# M82 should be visible left of center, M81 on the left hand edge. Two features caused by detector amp glow are also visible at the bottom of the image.
# ```

# ### Problem 1.2
#
# Finish the syntax for creating the `PSFPhot` class below. The init for the class should ask for three string paths: one to an image, one to a dark file, and one to a flat file. Within `__init__()`, use your `load_fits()` function to get the just the science image and header into the class. Store them as class attributes `self.data_init` and `self.data_header`.
#
# Don't forget to add a docstring to your init! It doesn't return anything, but sets the initial attributes.

# In[ ]:

class PSFPhot():

# ### Problem 1.3
#
# Add a method to your class called `dark_subtract()` which takes an image and a string-path to a dark file. Use `load_fits()` to read it in the dark, and subtract it from the image, and return the output.
#
# Add another method called `flat_field()` which takes in an image and a string-path to a flatfield file. Use `load_fits()` to read it in, normalize it to its maximum value, and divide the input image by it.
#
# Finally, within your init function, set a new class attribute called `self.data_calibrated`, and set it equal to the output of your calibration methods as run on the set of image paths that have been input.
#
# ```{hint}
# This step can be done in one line with nested functions if you prefer.
# ``
#
# Finally Use `implot()` to plot both the `data_init` and `data_calibrated` and show that your calibrations worked. You should see the amp glow regions dissapear, and the image sky should look nice and uniform.

# In[527]:

class PSFPhot():

# ```{note}
# Running the dark/flat methods within init means they'll happen right away when someone instantiates an object of this class. It's a matter of code design and personal preference to do this, versus having the user run those methods manually as part of the UX. I choose to wrap it into the init because it's a basic, common thing all inputs will need done to them.
# ```

# ### Problem 1.4
#
# The final step in setting up our image for proper analysis is the subtraction of the globally varying sky background. Performing that fit (generally a low order 2D polynomial of some kind) is beyond the scope of this lab, but we encourage you to look into this step more if you are interested.
#
# Instead, we're going to return to the `sep` package and take advantage of its background estimation feature. Recall it reads in the base image array and a separate mask. We want the user of this class to be able to supply a mask, so we won't run this step automatically in `__init__()`.
#
# Using the same stucture as in Lab 2, we'll calculate the background using `sep`.
#
# Add a method to your class that is called `subtract_background`, whose only argument (besides `self`) should be an optional mask (default `None`).
#
# Inside, use `sep` to measure the background and access, e.g., the `bkg.back()` method, and subtract this background-image from `self.data_calibrated`. You should set `mask=mask` in the sep call so your method input gets passed through. For convenient access, store the output of `bkg.back()` in a class attribute called `self.background`, and then set `self.image` equal to the background subtracted image. This method is known as a *setter*, because we're setting a class attribute but not returning anything.
#
# It can be handy here to add a print statement at the end of the method saying something like `"Background estimated; output saved to attribute 'image' "`. Since this is the main image we'll be using from now on, I've elected to give it the short name `image`. You may rather use something like `data_bg_sub`, for example
#
# ```{warning}
# Don't forget about the C order switch with sep.
# ```

# In[528]:

class PSFPhot():

# Our class usage will now look something like
# ```
# pipe = PSFPhot(im_path,dark_path,flat_path)
# pipe.subtract_background()
# ```
#
# or, if we create a mask (as we will below),
# ```
# pipe = PSFPhot(im_path,dark_path,flat_path)
# mask = # Something here
# pipe.subtract_background(mask=mask)
# ```
#

# ## Problem 2
#
# Our goal is to estimate the PSF of the above image, then measure fluxes of the stars and galaxies here accounting for the PSF.
#
# To start this process, we need to locate the stars in this image. We saw how to segment an image using `sep` last time, but in this lab we are going to carry out this step ourselves, using two methods.
#
# ### Problem 2.1
#
# Before we do this, we want to take the step of masking out several regions of the image which may register as peaks but which are not nicely isolated stars. In particular, the two galaxies need to be masked out when both estimating the background and when looking for point sources.
#
#
# Create a mask of the dimensions of the image data, containing `False` everywhere except where we want to mask out the galaxies (rough rectangular regions are fine for this masking). It's easiest to make the full array `False` first, then set the regions we want to mask to `True` via image indexing.
#
# Add a method to your class called `set_image_mask()` which overwrites `self.image` with a numpy masked array containing what used to be `self.image` as the data, and a mask that the user inputs.
#
# Finally, instantiate your class into a variable and check that the masking and subtraction are both working. You should plot your image by accessing it (it should be a class attribute by now). It should look like mine below.
# ```{warning}
# If we have the memory to, it's often worth avoiding overwriting class attributes. Here, we could set the masked image as `self.image_masked`, for example. I've elected not to here for two reasons. One: simplicity. Two: the unmasked version of the image is easily accessible still via `self.image.data`, the data attribute of numpy masked arrays.
# ```

# In[535]:

# Solution
mask = #define mask

class PSFPhot():

# In[536]:

# instantiate your class, run the subtraction, input the mask, see that it all works.

#
# ```{note}
# In our use case for this lab, the mask we will use when background subtracting is the same as the one we'll use for blocking out regions from being included in peak finding and such. But our code is flexible enough to, e.g., mask when background-subtracting but not when peak finding, or vice versa. Our mask setter only affects the image data, and the mask argument in the subtract background method only affects background subtraction.
# ```
#
# ```{tip}
# As it stands, `set_image_mask()` can only be run *after* `subtract_background()`, because that method sets `self.image` for the first time. We can use the `hasattr(self,'image')` check to see, when `set_image_mask()` is run, if that previous method was already run and `self.image` exists. For now, it's ok to assume the user will use the class in order, but a more robust code would either raise an exception or otherwise handle this case in some way.
# ```

# Plot the `.background` attribute of your class. It seems our image has a spatially varying background that includes a gradient across the image.

# In[ ]:

#plot the background

# ### Problem 2.2
#
# Now that we have the appropriate image regions masked, we can move on to the peak finder.
#
# The "fast" or "efficient" method of doing this involves some scipy filtering operations. But for our purposes, the "slow way" (iterating over the image pixels) takes ~few seconds to run, and is worth doing to build intuition.
#
# Add a method to your class to find peaks in an image by looping over each pixel and checking its neighbors, with a "peak" being defined as a region of higher flux than all adjacent pixels (i.e., the 8 surrounding pixels). In order to not pick up random noise pixels, also take an input called `threshold`. Within your algorithm, don't return any pixels which are "peaks" but for which the pixel value is below this threshold.
#
# ```{hint}
# :class: dropdown
# This algorithm needs to avoid the edges of the image (since indexing "i+1" won't work there). Creating a 1 pixel "buffer" in your looping will prevent this.
# ```

# In[70]:

# your code here

# The looping solution is slow, and will not scale well if we have to run on many images, but for one image is okay.
#
# There are several solutions which generally involve either **filtering** the image or **cross correlating** the image with a template. Here's one such solution.

# In[399]:

from scipy.ndimage import maximum_filter

def findpeaks_maxfilter(image, threshold):
    '''
    Algorithm for finding peaks (above a threshold) in an image

    Parameters
    ----------
    image: array_like
        2D array containing the image of interest.
    threshold: float
        minimum pixel value for inclusion in search

    Returns
    -------
    peaks: array_like
        array containing the x and y coordinates of peak regions.
    '''
    neighborhood = np.ones((3,3),dtype=bool) # just 3x3 True, defining the neighborhood over which to filter
    # find local maximum for each pixel
    amax = maximum_filter(image, footprint=neighborhood) #max filter will set each 9-square region in the image to the max in that region.

    peaks = np.where((image == amax) & (image >= threshold)) #find the pixels unaffected by the max filter.
    peaks = np.array([peaks[0],peaks[1]]).T
    return peaks

# Let's take a moment to understand how this algorithm works. The key is in the `maximum_filter()` step. Filtering a 2D image is a process carried out in fourier space, which is what allows scipy to carry it out quickly. But what is maximum filtering?
#
# ```{admonition} Definition
# Maximum Filtering is the process by which all pixels in local neighborhoods within an array are raised to the maximum value of any pixel in that neighborhood
# ```
#
# Let's look at a 1D case. Below, I define a 1D array that has some peaks in it.

# In[110]:

array_1dpeaks = np.array([1,1,2,1,2,1,2,1,2,3,4,3,2,1,2,4,5,6,8,6,5,4,3,2,4,3,4,2,1,0.5,1,2,1,2,3,2])

# Our data looks like this:

# In[122]:

fig, ax = plt.subplots(figsize=(6,6))
ax.plot(array_1dpeaks,lw=1.5,alpha=0.9)
ax.plot(array_1dpeaks,'o',color='C0')

# Let's now run the maximum filter on this data and plot it's output. I'm going to pick a neighborhood of 3, which means +/- 1 pixel around each location.

# In[112]:

mf = maximum_filter(array_1dpeaks, footprint=np.ones(3,dtype=bool))

# In[120]:

fig, ax = plt.subplots(figsize=(6,6))
ax.plot(array_1dpeaks,lw=1.5,alpha=0.9,label='original array')
ax.plot(array_1dpeaks,'o',color='C0')
ax.plot(mf,label='max filtered array')
eq, = np.where(array_1dpeaks==mf)
ax.plot(np.arange(len(mf))[eq],mf[eq],'o',label='local peaks')
ax.legend();

# What the filtering has done is for every 3 pixel neighborhood across this array, it's raised the value of all three pixels to the maximum value across the three. So we see that anywhere the three pixels were, e.g., (1,2,1), they are all now 2. What you should notice looking at this plot is that the max filtering has also identified true peaks in our data! Notice that the only spots where the orange curve (the max filtered version of the data) is equal to the original array is exactly at locations that are local maxima. This is because when applying max filtering to an array, the only values *unchanged* by the filtering are those that *are* local maxima (in the neighborhood defined).
#
# And thus, we have our peaks! All we need to do is find out `where()` the max filtered array equals the original array. Of course, we can also put in a threshold (in this example, maybe 2.5) to ensure low level noise doesn't enter in. This is why the `findpeaks_maxfilter()` function has a threshold option as well.
#
# ```{note}
# You may notice that the first index in the array is marked as a peak, despite not being one. Edges are always troublesome with these algorithms, and they normally have multiple options for how edges are handled, e.g, interpolating a constant value, reflecting over the edge, etc. For our lab, we're not going to worry about this.
# ```
#

# ### Problem 2.3
#
# How fast is max filtering over your looping algorithm?
#
# Use the `%%timeit` magic command in your notebook to test how fast the two algorithms are respectively. If you were working with a sample of 1000 images, how long would the looping algorithm take compared to the max-filtering case?
#
# ```{hint}
# Check out the "Timing Code" page in the "Quick Tips" sidebar tab of our website to see examples of using the `timeit` module.
# ```

# In[ ]:

# time your function

# In[ ]:

# time the max filter function

# How many times faster is the max filtering?

# In[ ]:

# Calculate

# How long would it take to run each on 1000 images (extrapolating from your results)?

# In[ ]:

# Calculate

# For the rest of the lab, you can either leave in your peak finder, or replace it with the max-filtering version (turning the function above into a method).
#
# Either way, add a peak-finding method into your class, which should be runnable with no inputs, but has an optional threshold argument to set a minimum flux value for inclusion.

# In[ ]:

# your class here

# ### Problem 2.4
#
# Run your peakfinder on your masked image, and assemble the list of peaks you'll use moving forward. Your peakfinder method should save a class attribute `self.peak_locations` which contains the ($x,y$) pairs of points associated with your peaks.
#
#
# Plot up the original image, but use our "source circling" technique from last lab to circle all the peaks found in the masked image. I show mine below.
#
# ```{hint}
# Recall that `implot()` returns a set of `fig, ax`, so you can then plot the circles onto the same axis as the image.
# ```

# In[549]:

# re-instantiate your object, run the subtraction, then your new peak finder here.

# ### Problem 2.5
#
# This should look pretty good -- most of what's circled above is clearly a star/point source in the image. However, one problem with this method is that single hot pixels are going to be registered as peaks, even via the clever algorithm. We need a way to eliminate these from our sample before moving on.
#
# Adjust your peak-finder method to add in something that checks that not only is there a true peak, but that at least 4 of the pixels around the peak are also elevated in flux (I used 0.5 times the peak flux). The easiest way is to loop over the peaks after they're found and institute the check --- there are far fewer peaks than pixels, so this doesn't significantly affect the runtime. But feel free to find a cleverer solution!
#
# ```{warning}
# Be careful with transpositions in this problem. when you plot coordinates, you plot(x,y), but when you index the image, you index image[y,x]. Always be tracking which is which!
# ```

# In[619]:

class PSFPhot():

# Re-find your peaks using your newer, better algorithm, and plot them below as before.

# In[ ]:

# your code here

# Notice that we've decreased our total number of peaks. But you should find that now, everything currently circled looks like a bright, "resolved" star. (resolved insofar as the PSF is spreading the light of the star over multiple pixels).
#
# ### Problem 2.6
#
# In your the image above, you should see that ~8-10 stars look like they are circled by several very closely overlapping circles all targeting the same star. Infer (or investigate and determine) why this has happened, and write your answer below.
#
#
#

# *answer here*

# ## Problem 3
#
# We now have a function that can return the peaks in a given image. Our next step is going to be to estimate the exact center of those peaks (stars) using their **centroid**.
#
# ```{admonition} Definition
# The centroid is the light-weighted-mean of a set of pixels. It is not always the maximum-valued pixel, and is determined to sub-pixel accuracy.
# ```
#
# Many of you have seen the centroid formula, but as a reminder, it looks like this (in 1D):
#
#
# $$
# x_{\rm com} = \frac{\sum{x_i \hat{f}_i}}{\sum \hat{f}_i},
# $$
#
# where $x_i$ are the positions and $\hat{f}_i$ are the fluxes at those positions.
#
# In 2D, when working with images, the $x$ and $y$ centers of mass are independent, and the 2D centroid is just the location ($x_{\rm com}$, $y_{\rm com}$).
#
# ### Problem 3.1
#
# Add a method to your class called `centroid()` which should read in an $x$ and $y$ peak location (these will be returned by your peak finder), then creates a cutout by indexing `self.image` in a window of N (user-settable), and determine the centroid of this window. The $x,y$ location of this centroid should be returned.
#
# One subtlety --- We want to use a window size greater than 1 pixel on either side of each peak. But because our peak finder is likely to have found peaks near the edge of the detector (both because it needs only 1 pixel-thick borders and because it handles edges), if we write a centroid function that, e.g., uses a 10x10 pixel window, we'll end up trying to index over the edge of the original image whenever there's a star near the edge. Because of this, your function should raise an exception if a peak position is entered whose distance from an edge is less than half the window size.

# In[622]:

class PSFPhot():
    # copy down and add to your class

# ### Problem 3.2
# Use your `centroid()` function to confirm that the algorithm is working by testing it on a few individual peaks from your peak list, and make a plot showing the window region and the determined centroid (along with the location of the input peak). I'm leaving behind a demo of what I mean below. The blue point is the pixel with the peak flux, while the crosshairs center on the determined centroid
#
# ```{Note}
# It's usually ok if the centroid is not at what appears to be the very center of the light distribution of the star. Often due to seeing and detector effects, along with tricks of the stretch you happen to be using, the centroid doesn't look perfectly centered. But it shouldn't be super off, either -- see below.
# ```
# In my solution below, I plot 25 stars from my peak finder, with peaks and centroids shown. You should get similar looking results as you query stars in frame.
#

# In[624]:

# You need not plot 25, but check a handful of stars to make sure you're in business.

# ### Problem 3.3
#
# If you recall from above, you determined why the peak algorithm occasionally marked a bunch of pixels as peaks within the same star, which shouldn't happen. It should be clear from your answer that these stars will be unusable for the purposes of measuring the PSF of stars in the image. We thus need to remove these ~8-10 stars from our sample.
#
# Write a method `check_stars()` which takes in the list of centroids, and identifies these cases. The easiest way to do this is to simply iterate through the list of centroids and compute the distance to all other centroids in the list. Any subgroups with small distances should be removed. (Say, 5-10 pixels).
#
#
# This method should return the final list of centroids for use. Plot them over the data to confirm these "stacked" peak cases have been removed.
#

# In[ ]:

# your code here

# ## Problem 4
#
# Armed with a dark-subtracted, flat-fielded, background-subtracted image, as well as with a list of centroids corresponding to stars in our image, we are ready to estimate the PSF.
#
# There are two main functional forms typically used to fit star profiles: 2D Gaussians, and Moffat profiles (which combines the shapes of a Gaussian and Lorentzian to best match both the inner and outer regions of the PSF).
#
# We're going to use the [`Gaussian2D`](https://docs.astropy.org/en/stable/api/astropy.modeling.functional_models.Gaussian2D.html) class from `astropy` to do this:

# In[433]:

from astropy.modeling.functional_models import Gaussian2D

# For each star, a Gaussian2D profile (normalized) will be used "as the PSF". The parameters we need to know for this profile are $x,y$, for which we'll use the centroids we calculated earlier, the amplitude (set by the normalization), and $\sigma_x,\sigma_y$, the standard deviations in the two axes. For this lab, we're going to assume our stars are circular ($\sigma_x=\sigma_y$). This is a strictly incorrect, but not a bad assumption for most cases. All other optional arguments we won't need, primarily due to the assumption of circularity.
#
# ```{note}
# We are going to make a point estimate of the "size" of the stars in our image, which constrains us from using a more fancy model for the PSF. An example of a more sophisticated setup would be *fitting* a Gaussian or Moffat profile to every star, and in a Bayesian framework marginalizing over the stars to determine the best-fit PSF (including ellipticity, etc) for the image, or, even fancier, interpolating a PSF model which varies over the detector.
# ```
#
# PSF photometry works by multiplying the *data* (say, a cutout around a star) by the *estimated PSF* during the fluxing stage. Instead of picking a radius and performing aperture photometry (which includes fully all pixels within the aperture and throws out all pixels beyond), this method attempts to weight each pixel fractionally by how likely it is to be stellar flux, with the weighting coming from the PSF of the detector. This means further pixels may still be included, but will contribute less than pixels near the center of the star.
#
# The formula for measuring the PSF flux of a star is
#
# $$
# f_{\rm PSF} = \frac{\sum \hat{f_i} p_i}{\sum p_i^2},
# $$
#
# where $\hat{f_i}$ are the fluxes in your image and $p_i$ is your PSF estimate. This formula should be reminiscent of the centroiding formula; it's a similar weighting scheme.

# `Gaussian2D` is a class, but we want to interact with it pretty simply, and have simplified inputs. I've made a quick wrapper function below which allows us to enter $\sigma_x$ and $\sigma_y$ and then $x,y$ grids created via `np.meshgrid()`, and creates the Gaussian and evaluates it on our grid.

# In[490]:

def eval_gauss(x_arr,y_arr,sigma_x,sigma_y,mu_x,mu_y):

    g = Gaussian2D.evaluate(x=x_arr,y=y_arr,amplitude=1,theta=0,x_mean=mu_x,
                   y_mean=mu_y,
                   x_stddev=sigma_x,
                   y_stddev=sigma_y)
    g/=np.sum(g)
    return g

# In[491]:

xx, yy = np.meshgrid(np.arange(loc[0]-10,loc[0]+10),
                     np.arange(loc[1]-10,loc[1]+10))
model = eval_gauss(x_arr=xx,y_arr=yy,sigma=3,mu_x=testx,mu_y=testy)

# In[485]:

fig, ax = plt.subplots(figsize=(8,8))
ax.imshow(model,origin='lower')

# As we can see, I now have a model for the PSF which I can easily create for given inputs. We're going to do this for cutouts around each star, and instead of a random $\sigma$, we're going to estimated it using the second moment (moment of inertia) of the star itself.
#
# The formula for this (from Markevich et al. 1989) is
#
# $$
# \sigma_x = \left[\frac{\sum x_i^2 \hat{f}_i}{\sum \hat{f}_i} - x_{\rm com}^2\right]^{1/2}
# $$
#
# $$
# \sigma_y = \left[\frac{\sum y_i^2 \hat{f}_i}{\sum \hat{f}_i} - y_{\rm com}^2\right]^{1/2}
# $$
#
# In this case, we'll need to use `meshgrid()` directly within our second moment function, as you can see it depends on the difference between the pixels and the centroid.
#
# ### Problem 4.1
#
# Add a method to your class called `second_moment` which reads in an image cutout, the meshgrid input (xx,yy) (or constructs it), and finally, a centroid (x and y). Inside, use the formulas above to determine sigma_x and sigma_y, and return them

# In[494]:

# Below, I show that a functional form of my second moment code returns 2.92 and 2.98 when run on the example psf I made above (which had a true sigma of 3).

# In[495]:

second_moment(model,xx,yy,testx,testy)

# Within 10%, this calculation, which requires no fitting, tells us $\sigma$ (under the assumption the distribution is Gaussian). Thus, by running out image cutouts through this function, we can derive for ourselves a good $\sigma$ to choose in our Gaussian model of the PSF.

# ### Problem 4.2
#
# We now have everything we need to run our pipeline.
#
# Right now, our final three methods (centroider, star checker, and width-finder) all *return* their values. I had you do it that way for testing purposes. To finalize our "object oriented pipeline" though, we should make a wrapper method that will bring them together and handle the outputs.
#
#
# Add a new, final method called `psf_photometry`. When the user runs this method, it should first feed the peaks into the centroid code one by one, assembling a set of centroids. It should then construct cutouts of `self.image` around each peak (or centroid), and feed those, plus the centroids, and appropriate meshgrids into the second moment function to save a pair sigma_x and sigma_y for each star as well. And finally, it should use the eval_gauss function I've provided above to create a reasonable PSF model, then carry out the PSF photometry (using the equation above).
#
# At the end of it all, you can save the centroids, widths, and psf-fluxes all to class attributes, where they will remain accessible. You can also find a nice way to return them for extra credit (see below).

# In[ ]:

# I know this was long! Congrats for getting here!

# ## Bonus Problem (+1 Extra Credit)
#
# A convenient way to return the set of measurements carried out by our pipeline is via a `flux table`, like the one returned by `aperture_photometry` from `astropy`. For our purposes, since we'll be spending a lot of time with data frames, let's go with that.
#
# Modify your final `psf_photometry` method to assemble a nice `DataFrame` which contains as colummns `peak_x`, `peak_y`, `centroid_x`, `centroid_y`, `width_x`, `width_y`, and `psf_flux`. Each row should be the relevant values for a given star that was originally found by your peakfinder (but which made it through our sample cleaning stage).
#
# ```{tip}
# Inside your final method, you probably have lists or arrays containing the output of your centroid and width methods as they looped over the stars (and your find peak function returns a list of positions outright). So if you set up an empty DataFrame, you can set, e.g., `df['peak_x'] = peaks[:,0]` or something like `df['width_x'] = out_widths`, to set up these columns. As long as you never messed with the order of the lists, it should turn into the desired output frame.
# ```

# In[ ]: