# -*- coding: utf-8 -*-
"""
Created on Sat Aug  5 19:48:34 2017

@author: max
"""

import numpy as np
import matplotlib.pyplot as plt
from math import pi, sqrt
from collections import namedtuple
from .Functions.CalcIOLMaxwellian import calc_iol_maxwellian
from .Functions.CalcIOLMonoEn import calc_iol_mono_en
from .Functions.CalcIOLBeams import calc_iol_beams
from GT3.utilities.PlotBase import PlotBase
import sys
import GT3.constants as constants
from GT3 import Core

m_d = constants.deuteron_mass
m_t = constants.triton_mass
m_c = constants.carbon_mass
m_a = constants.alpha_mass

class IOL(PlotBase):
    """
    """

    def __init__(self, inp, core: Core):
        super().__init__()
        sys.dont_write_bytecode = True
        np.warnings.filterwarnings('ignore')

        numcos = inp.numcos  # TODO: Clean this up
        angles = np.linspace(-1, 1, inp.numcos + 1)
        self.coslist = ((angles + np.roll(angles, -1)) / 2)[:-1]
        self.rho = core.rho
        self.set_plot_rho1d(self.rho[:, 0])
        self.calc_iol_beams = calc_iol_beams
        polpts = len(core.rho[-1])
        radpts = len(core.rho.T[-1])

        # THE FOLLOWING ARRAYS ARE 4-DIMENSIONAL ARRAYS
        # [ LAUNCH THETA POSITION , LAUNCH ANGLE COSINE, LAUNCH r  , EXIT THETA POSITION  ]

        # NOTE TO FUTURE DEVELOPERS: IF YOU TRY TO LOOP OVER THE PLASMA POINTS, LAUNCH ANGLES, AND
        # EXIT LOCATIONS IN PYTHON THE WAY YOU WOULD IN C OR FORTRAN, IT'S GOING TO TAKE FOREVER.
        # ALTHOUGH THESE ARRAYS TAKE MORE MEMORY THAT I'D LIKE, IT'S CURRENTLY NECESSARY TO DO IT THIS WAY.
        # MAYBE SOMETHING TO IMPROVE ON IN THE FUTURE.

        # CREATE ARRAYS FOR LAUNCH POINTS IN THE PLASMA
        r0 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                 np.ones(inp.numcos)[:, None, None],
                                 core.rho)[-1]

        B0 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                 np.ones(inp.numcos)[:, None, None],
                                 core.B.tot.val)[-1]

        f0 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                 np.ones(inp.numcos)[:, None, None],
                                 core.f_phi)[-1]

        psi0 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                   np.ones(inp.numcos)[:, None, None],
                                   core.psi.psi)[-1]

        phi0 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                   np.ones(inp.numcos)[:, None, None],
                                   core.E_pot.val)[-1]  # * 1E3  # now in volts

        zeta0 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                    self.coslist[:, None, None],
                                    np.ones(core.R.shape))[1]

        # CREATE ARRAYS FOR DESTINATION POINTS ALONG THE SEPERATRIX
        R1 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                 np.ones(inp.numcos)[:, None, None],
                                 np.ones(radpts)[:, None],
                                 np.ones(polpts)[:],
                                 core.R[-1][:, None, None, None])[-1]

        f1 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                 np.ones(inp.numcos)[:, None, None],
                                 np.ones(radpts)[:, None], np.ones(polpts)[:],
                                 core.f_phi[-1][:, None, None, None])[-1]

        B1 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                 np.ones(inp.numcos)[:, None, None],
                                 np.ones(radpts)[:, None],
                                 np.ones(polpts)[:],
                                 core.B.tot.val[-1][:, None, None, None])[-1]

        psi1 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                   np.ones(inp.numcos)[:, None, None],
                                   np.ones(radpts)[:, None],
                                   np.ones(polpts)[:],
                                   core.psi.psi[-1][:, None, None, None])[-1]

        phi1 = np.broadcast_arrays(np.ones(polpts)[:, None, None, None],
                                   np.ones(inp.numcos)[:, None, None],
                                   np.ones(radpts)[:, None],
                                   np.ones(polpts)[:],
                                   core.E_pot.val[-1][:, None, None, None])[-1]  # * 1E3  # now in volts

        Tprofile = namedtuple('Tprofile', 'i C')(
            core.T.i.kev.val.T[0],
            core.T.C.kev.val.T[0]
        )

        # iol_params = {}
        # iol_params['r0'] = r0
        # iol_params['B0'] = B0
        # iol_params['f0'] = f0
        # iol_params['psi0'] = psi0
        # iol_params['phi0'] = phi0
        # iol_params['zeta0'] = zeta0
        # iol_params['R1'] = R1
        # iol_params['f1'] = f1
        # iol_params['B1'] = B1
        # iol_params['psi1'] = psi1
        # iol_params['phi1'] = phi1
        # iol_params['Tprofile'] = Tprofile

        # convert iol_params to namedtuple so individual parameters can be accessed as normal attributes
        self.iol_p = namedtuple('iol_p', 'r0 B0 f0 psi0 phi0 zeta0 R1 f1 B1 psi1 phi1')(
            r0,
            B0,
            f0,
            psi0,
            phi0,
            zeta0,
            R1,
            f1,
            B1,
            psi1,
            phi1
        )

        # Calculate IOL for thermal deuterium
        forb_d_therm, morb_d_therm, eorb_d_therm = calc_iol_maxwellian(1,
                                                                       m_d,
                                                                       self.iol_p,
                                                                       core.thetapts,
                                                                       Tprofile.i,
                                                                       self.coslist,
                                                                       numcos)
        self.forb_d_therm = inp.R_loss * forb_d_therm
        self.morb_d_therm = inp.R_loss * morb_d_therm
        self.eorb_d_therm = inp.R_loss * eorb_d_therm
        self.forb_d_therm_1D = self.forb_d_therm[:, 0]
        self.morb_d_therm_1D = self.morb_d_therm[:, 0]
        self.eorb_d_therm_1D = self.eorb_d_therm[:, 0]

        # Calculate IOL for thermal tritium
        forb_t_therm, morb_t_therm, eorb_t_therm = calc_iol_maxwellian(1,
                                                                       m_t,
                                                                       self.iol_p,
                                                                       core.thetapts,
                                                                       Tprofile.i,
                                                                       self.coslist,
                                                                       numcos)
        self.forb_t_therm = inp.R_loss * forb_t_therm
        self.morb_t_therm = inp.R_loss * morb_t_therm
        self.eorb_t_therm = inp.R_loss * eorb_t_therm
        self.forb_t_therm_1D = self.forb_t_therm[:, 0]
        self.morb_t_therm_1D = self.morb_t_therm[:, 0]
        self.eorb_t_therm_1D = self.eorb_t_therm[:, 0]

        # Calculate IOL for thermal carbon
        forb_c_therm, morb_c_therm, eorb_c_therm = calc_iol_maxwellian(6,
                                                                       m_c,
                                                                       self.iol_p,
                                                                       core.thetapts,
                                                                       Tprofile.C,
                                                                       self.coslist,
                                                                       numcos)
        self.forb_c_therm = inp.R_loss * forb_c_therm
        self.morb_c_therm = inp.R_loss * morb_c_therm
        self.eorb_c_therm = inp.R_loss * eorb_c_therm
        self.forb_c_therm_1D = self.forb_c_therm[:, 0]
        self.morb_c_therm_1D = self.morb_c_therm[:, 0]
        self.eorb_c_therm_1D = self.eorb_c_therm[:, 0]

        # Calculate IOL for thermal alphas
        forb_a_therm, morb_a_therm, eorb_a_therm = calc_iol_maxwellian(2,
                                                                       m_a,
                                                                       self.iol_p,
                                                                       core.thetapts,
                                                                       Tprofile.i,
                                                                       self.coslist,
                                                                       numcos)
        self.forb_a_therm = inp.R_loss * forb_a_therm
        self.morb_a_therm = inp.R_loss * morb_a_therm
        self.eorb_a_therm = inp.R_loss * eorb_a_therm
        self.forb_a_therm_1D = self.forb_a_therm[:, 0]
        self.morb_a_therm_1D = self.morb_a_therm[:, 0]
        self.eorb_a_therm_1D = self.eorb_a_therm[:, 0]

        # Calculate IOL for fast, monoenergetic alphas
        v_alpha = sqrt(2 * 3.5E6 * 1.6021E-19 / m_a)
        forb_a_fast, morb_a_fast, eorb_a_fast = calc_iol_mono_en(2,
                                                                 m_a,
                                                                 self.iol_p,
                                                                 core.thetapts,
                                                                 v_alpha,
                                                                 self.coslist,
                                                                 numcos)

        # currently applying R_loss to fast alphas as well as thermal, although I'm skeptical of this. -MH
        self.forb_a_fast = inp.R_loss * forb_a_fast
        self.morb_a_fast = inp.R_loss * morb_a_fast
        self.eorb_a_fast = inp.R_loss * eorb_a_fast
        self.forb_a_fast_1D = self.forb_a_fast[:, 0]
        self.morb_a_fast_1D = self.morb_a_fast[:, 0]
        self.eorb_a_fast_1D = self.eorb_a_fast[:, 0]

        # Calculate IOL for neutral deuterium beams
        v_beam = sqrt(2 * 80.0E3 * 1.6021E-19 / m_d)
        zeta_beam = -0.96
        forb_d_nbi, morb_d_nbi, eorb_d_nbi = self.calc_iol_beams(1,
                                                                 m_d,
                                                                 self.iol_p,
                                                                 core.thetapts,
                                                                 v_beam,
                                                                 zeta_beam,
                                                                 self.coslist)
        # currently applying R_loss to fast alphas as well as thermal, although I'm skeptical of this. -MH
        self.forb_d_nbi = inp.R_loss * forb_d_nbi
        self.morb_d_nbi = inp.R_loss * morb_d_nbi
        self.eorb_d_nbi = inp.R_loss * eorb_d_nbi
        self.forb_d_nbi_1D = self.forb_d_nbi[:, 0]
        self.morb_d_nbi_1D = self.morb_d_nbi[:, 0]
        self.eorb_d_nbi_1D = self.eorb_d_nbi[:, 0]

    def plot_F_i(self, edge=True):
        """
        Plots the 1D GT3.IOL differential ion number loss fraction
        """
        fig = self._plot_base(self.forb_d_therm_1D, yLabel=r'$F(\rho)$',
                              title="GT3.IOL differential ion number loss fraction", edge=edge, color='blue')
        return fig

    def plot_M_i(self, edge=True):
        """
        Plots the 1D GT3.IOL differential ion momentum loss fraction
        """
        fig = self._plot_base(self.morb_d_therm_1D, yLabel=r'$M(\rho)$',
                              title="GT3.IOL differential ion momentum loss fraction", edge=edge, color='green')
        return fig

    def plot_E_i(self, edge=True):
        """
        Plots the 1D GT3.IOL differential ion energy loss fraction
        """
        fig = self._plot_base(self.eorb_d_therm_1D, yLabel=r'$E(\rho)$',
                              title="GT3.IOL differential ion energy loss fraction", edge=edge, color='red')
        return fig

    def plot_all_i(self):
        plot = plt.figure()
        fig1 = plot.add_subplot(111)
        fig1.set_xlabel(r'$\rho$', fontsize=20)
        fig1.set_ylabel(r'$\frac{\partial F}{\partial r}$', fontsize=25)
        fig1.set_title('Ion number loss fraction')
        fig1.scatter(self.rho, self.forb_d_therm_1D, marker='o', color='blue')

        fig2 = plot.add_subplot(121)
        fig2.set_xlabel(r'$\rho$', fontsize=20)
        fig2.set_ylabel(r'$\frac{\partial E}{\partial r}$', fontsize=25)
        fig2.set_title('Ion energy loss fraction')
        fig2.scatter(self.rho, self.eorb_d_therm_1D, marker='o', color='red')

        fig3 = plot.add_subplot(131)
        fig3.set_xlabel(r'$\rho$', fontsize=20)
        fig3.set_ylabel(r'$\frac{\partial M}{\partial r}$', fontsize=25)
        fig3.set_title('Ion momentum loss fraction')
        fig3.scatter(self.rho, self.morb_d_therm_1D, marker='o', color='green')

        return plot

    def plot_F_C(self, edge=True):
        """
        Plots the 1D GT3.IOL differential carbon number loss fraction
        """
        fig = self._plot_base(self.forb_c_therm_1D, yLabel=r'$\frac{\partial F}{\partial r}$',
                              title="GT3.IOL differential carbon number loss fraction", edge=edge, color='blue')
        return fig

    def plot_M_C(self, edge=True):
        """
        Plots the 1D GT3.IOL differential carbon momentum loss fraction
        """
        fig = self._plot_base(self.morb_c_therm_1D, yLabel=r'$\frac{\partial M}{\partial r}$',
                              title="GT3.IOL differential carbon momentum loss fraction", edge=edge, color='green')
        return fig

    def plot_E_C(self, edge=True):
        """
        Plots the 1D GT3.IOL differential carbon energy loss fraction
        """
        fig = self._plot_base(self.eorb_c_therm_1D, yLabel=r'$\frac{\partial E}{\partial r}$',
                              title="GT3.IOL differential carbon energy loss fraction", edge=edge, color='green')
        return fig

    def plot_all_C(self):
        plot = plt.figure()
        fig1 = plot.add_subplot(111)
        fig1.set_xlabel(r'$\rho$', fontsize=20)
        fig1.set_ylabel(r'$\frac{\partial F}{\partial r}$', fontsize=25)
        fig1.set_title('Carbon number loss fraction')
        fig1.scatter(self.rho, self.forb_c_therm_1D, marker='o', color='blue')

        fig2 = plot.add_subplot(121)
        fig2.set_xlabel(r'$\rho$', fontsize=20)
        fig2.set_ylabel(r'$\frac{\partial E}{\partial r}$', fontsize=25)
        fig2.set_title('Carbon energy loss fraction')
        fig2.scatter(self.rho, self.eorb_c_therm_1D, marker='o', color='red')

        fig3 = plot.add_subplot(131)
        fig3.set_xlabel(r'$\rho$', fontsize=20)
        fig3.set_ylabel(r'$\frac{\partial M}{\partial r}$', fontsize=25)
        fig3.set_title('Carbon momentum loss fraction')
        fig3.scatter(self.rho, self.morb_c_therm_1D, marker='o', color='green')

        return plot

    def plot_F_i_fast(self, edge=True):
        """
        Plots the 1D GT3.IOL differential fast ion number loss fraction
        """
        fig = self._plot_base(self.forb_d_nbi_1D, yLabel=r'$\frac{\partial F}{\partial r}$',
                              title="GT3.IOL differential fast ion number loss fraction", edge=edge, color='blue')
        return fig

    def plot_M_i_fast(self,edge=True):
        """
        Plots the 1D GT3.IOL differential fast ion momentum loss fraction
        """
        fig = self._plot_base(self.morb_d_nbi_1D, yLabel=r'$\frac{\partial M}{\partial r}$',
                              title="GT3.IOL differential fast ion momentum loss fraction", edge=edge, color='green')
        return fig

    def plot_E_i_fast(self, edge=True):
        """
        Plots the 1D GT3.IOL differential fast ion energy loss fraction
        """
        fig = self._plot_base(self.eorb_d_nbi_1D, yLabel=r'$\frac{\partial E}{\partial r}$',
                              title="GT3.IOL differential fast ion energy loss fraction", edge=edge, color='red')
        return fig

    def plot_all_i_fast(self):
        plot = plt.figure()
        fig1 = plot.add_subplot(111)
        fig1.set_xlabel(r'$\rho$', fontsize=20)
        fig1.set_ylabel(r'$\frac{\partial F}{\partial r}$', fontsize=25)
        fig1.set_title('Fast ion number loss fraction')
        fig1.scatter(self.rho, self.forb_d_nbi_1D, marker='o', color='blue')

        fig2 = plot.add_subplot(121)
        fig2.set_xlabel(r'$\rho$', fontsize=20)
        fig2.set_ylabel(r'$\frac{\partial E}{\partial r}$', fontsize=25)
        fig2.set_title('Fast ion energy loss fraction')
        fig2.scatter(self.rho, self.eorb_d_nbi_1D, marker='o', color='red')

        fig3 = plot.add_subplot(131)
        fig3.set_xlabel(r'$\rho$', fontsize=20)
        fig3.set_ylabel(r'$\frac{\partial M}{\partial r}$', fontsize=25)
        fig3.set_title('Fast ion momentum loss fraction')
        fig3.scatter(self.rho, self.morb_d_nbi_1D, marker='o', color='green')

        return plot


import time
import torch
import math
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import utils
import numpy as np

def build_targets(pred_boxes, target, anchors, num_anchors, num_classes, nH, nW, noobject_scale, object_scale, sil_thresh, seen):
    nB = target.size(0)
    nA = num_anchors
    nC = num_classes
    anchor_step = len(anchors)//num_anchors
    conf_mask  = torch.ones(nB, nA, nH, nW) * noobject_scale
    conf_mask = conf_mask.view(nB, -1)
    coord_mask = torch.zeros(nB, nA, nH, nW)
    cls_mask   = torch.zeros(nB, nA, nH, nW)
    tx         = torch.zeros(nB, nA, nH, nW)
    ty         = torch.zeros(nB, nA, nH, nW)
    tw         = torch.zeros(nB, nA, nH, nW)
    th         = torch.zeros(nB, nA, nH, nW)
    tconf      = torch.zeros(nB, nA, nH, nW)
    tcls       = torch.zeros(nB, nA, nH, nW)

    nAnchors = nA*nH*nW
    nPixels  = nH*nW
    for b in range(nB):
        cur_pred_boxes = pred_boxes[b*nAnchors:(b+1)*nAnchors].t()
        cur_ious = torch.zeros(nAnchors)
        for t in range(50):
            if target[b][t*5+1] == 0:
                break
            gx = target[b][t*5+1]*nW
            gy = target[b][t*5+2]*nH
            gw = target[b][t*5+3]*nW
            gh = target[b][t*5+4]*nH
            cur_gt_boxes = torch.FloatTensor([gx,gy,gw,gh]).repeat(nAnchors,1).t()
            cur_ious = torch.max(cur_ious, utils.bbox_ious(cur_pred_boxes, cur_gt_boxes, x1y1x2y2=False))
        conf_mask[b][cur_ious>sil_thresh] = 0
    if seen < 12800:
       if anchor_step == 4:
           tx = torch.FloatTensor(anchors).view(nA, anchor_step).index_select(1, torch.LongTensor([2])).view(1,nA,1,1).repeat(nB,1,nH,nW)
           ty = torch.FloatTensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([2])).view(1,nA,1,1).repeat(nB,1,nH,nW)
       else:
           tx.fill_(0.5)
           ty.fill_(0.5)
       tw.zero_()
       th.zero_()
       coord_mask.fill_(1)

    conf_mask = conf_mask.view(nB, nA, nH, nW)

    nGT = 0
    nCorrect = 0
    for b in range(nB):
        for t in range(50):
            if target[b][t*5+1] == 0:
                break
            nGT = nGT + 1
            best_iou = 0.0
            best_n = -1
            min_dist = 10000
            gx = target[b][t*5+1] * nW
            gy = target[b][t*5+2] * nH
            gi = int(gx)
            gj = int(gy)
            gw = target[b][t*5+3]*nW
            gh = target[b][t*5+4]*nH
            gt_box = [0, 0, gw, gh]
            for n in range(nA):
                aw = anchors[anchor_step*n]
                ah = anchors[anchor_step*n+1]
                anchor_box = [0, 0, aw, ah]
                iou  = utils.bbox_iou(anchor_box, gt_box, x1y1x2y2=False)
                if anchor_step == 4:
                    ax = anchors[anchor_step*n+2]
                    ay = anchors[anchor_step*n+3]
                    dist = pow(((gi+ax) - gx), 2) + pow(((gj+ay) - gy), 2)
                if iou > best_iou:
                    best_iou = iou
                    best_n = n
                elif anchor_step==4 and iou == best_iou and dist < min_dist:
                    best_iou = iou
                    best_n = n
                    min_dist = dist

            gt_box = [gx.cuda(), gy.cuda(), gw.cuda(), gh.cuda()]
            pred_box = pred_boxes[b*nAnchors+best_n*nPixels+gj*nW+gi].cuda()

            coord_mask[b][best_n][gj][gi] = 1
            cls_mask[b][best_n][gj][gi] = 1
            conf_mask[b][best_n][gj][gi] = object_scale
            tx[b][best_n][gj][gi] = target[b][t*5+1] * nW - gi
            ty[b][best_n][gj][gi] = target[b][t*5+2] * nH - gj
            tw[b][best_n][gj][gi] = math.log(gw/anchors[anchor_step*best_n])
            th[b][best_n][gj][gi] = math.log(gh/anchors[anchor_step*best_n+1])
            iou = utils.bbox_ious(gt_box, pred_box, x1y1x2y2=False) # best_iou
            tconf[b][best_n][gj][gi] = iou
            tcls[b][best_n][gj][gi] = target[b][t*5]
            if iou > 0.5:
                nCorrect = nCorrect + 1

    return nGT, nCorrect, coord_mask, conf_mask, cls_mask, tx, ty, tw, th, tconf, tcls

class YoloLayer(nn.Module):
    def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1):
        super(YoloLayer, self).__init__()
        self.anchor_mask = anchor_mask
        self.num_classes = num_classes
        self.anchors = anchors
        self.num_anchors = num_anchors
        self.anchor_step = len(anchors)/num_anchors
        self.coord_scale = 1
        self.noobject_scale = 1
        self.object_scale = 5
        self.class_scale = 1
        self.thresh = 0.6
        self.stride = 32
        self.seen = 0

    def forward(self, output, target=None):
        if self.training:
            #output : BxAs*(4+1+num_classes)*H*W
            losses = []
#            for o_ind, output in enumerate(outputs):
            t0 = time.time()
            nB = output.size(0)
#            nA = self.num_anchors//3
            #nA = self.num_anchors
            nA = len(self.anchor_mask)
            nC = self.num_classes
            nH = output.size(2)
            nW = output.size(3)
            anchors = []
            for am in self.anchor_mask:
                anchors.append(self.anchors[2*am])
                anchors.append(self.anchors[2*am + 1])

            output   = output.view(nB, nA, (5+nC), nH, nW)
            x    = output.index_select(2, Variable(torch.cuda.LongTensor([0]))).view(nB, nA, nH, nW)
            x    = F.sigmoid(x)
            y    = output.index_select(2, Variable(torch.cuda.LongTensor([1]))).view(nB, nA, nH, nW)
            y    = F.sigmoid(y)
            width= output.index_select(2, Variable(torch.cuda.LongTensor([2]))).view(nB, nA, nH, nW) / 416
            height= output.index_select(2, Variable(torch.cuda.LongTensor([3]))).view(nB, nA, nH, nW) / 416
            conf = output.index_select(2, Variable(torch.cuda.LongTensor([4]))).view(nB, nA, nH, nW)
            conf = F.sigmoid(conf)
            cls  = output.index_select(2, Variable(torch.linspace(5,5+nC-1,nC)).long().cuda())
            cls  = cls.view(nB*nA, nC, nH*nW).transpose(1,2).contiguous().view(nB*nA*nH*nW, nC)
            t1 = time.time()

            pred_boxes = torch.cuda.FloatTensor(4, nB*nA*nH*nW)
            grid_x = torch.linspace(0, nW-1, nW).repeat(nH,1).repeat(nB*nA, 1, 1).view(nB*nA*nH*nW).cuda()
            grid_y = torch.linspace(0, nH-1, nH).repeat(nW,1).t().repeat(nB*nA, 1, 1).view(nB*nA*nH*nW).cuda()
#            anchor_w = torch.Tensor(self.anchors[o_ind*nA*2:(o_ind+1)*nA*2]).view(nA, self.anchor_step).index_select(1, torch.LongTensor([0])).cuda()
#            anchor_h = torch.Tensor(self.anchors[o_ind*nA*2:(o_ind+1)*nA*2]).view(nA, self.anchor_step).index_select(1, torch.LongTensor([1])).cuda()
#            anchor_w = torch.Tensor(self.anchors).view(nA, self.anchor_step).index_select(1, torch.LongTensor([0])).cuda()
#            anchor_h = torch.Tensor(self.anchors).view(nA, self.anchor_step).index_select(1, torch.LongTensor([1])).cuda()
            anchor_w = torch.Tensor(anchors).view(nA, self.anchor_step).index_select(1, torch.LongTensor([0])).cuda()
            anchor_h = torch.Tensor(anchors).view(nA, self.anchor_step).index_select(1, torch.LongTensor([1])).cuda()

            anchor_w = anchor_w.repeat(nB, 1).repeat(1, 1, nH*nW).view(nB*nA*nH*nW)
            anchor_h = anchor_h.repeat(nB, 1).repeat(1, 1, nH*nW).view(nB*nA*nH*nW)
            pred_boxes[0] = x.view(-1) + grid_x
            pred_boxes[1] = y.view(-1) + grid_y
            pred_boxes[2] = torch.exp(width).view(-1) * anchor_w
            pred_boxes[3] = torch.exp(height).view(-1) * anchor_h
            pred_boxes = utils.convert2cpu(pred_boxes.transpose(0,1).contiguous().view(-1,4))
            t2 = time.time()

            nGT, nCorrect, coord_mask, conf_mask, cls_mask,\
            tx, ty, tw, th, tconf, tcls = \
            build_targets(pred_boxes, target, anchors, # self.anchors
                          nA, nC, nH, nW,
                          self.noobject_scale, self.object_scale,
                          self.thresh, self.seen)
            cls_mask = (cls_mask == 1)
            nProposals = int((conf > 0.25).sum())

            tx    = Variable(tx.cuda())
            ty    = Variable(ty.cuda())
            tw    = Variable(tw.cuda())
            th    = Variable(th.cuda())
            tconf = Variable(tconf.cuda())
            tcls  = Variable(tcls[cls_mask].long().cuda())

            coord_mask = Variable(coord_mask.cuda())
            conf_mask  = Variable(conf_mask.cuda().sqrt())
            cls_mask   = Variable(cls_mask.view(-1, 1).repeat(1,nC).cuda())
            cls        = cls[cls_mask].view(-1, nC)

            t3 = time.time()

            loss_x = self.coord_scale * nn.MSELoss(size_average=True)(x*coord_mask, tx*coord_mask)/2.0
            loss_y = self.coord_scale * nn.MSELoss(size_average=True)(y*coord_mask, ty*coord_mask)/2.0
            loss_w = self.coord_scale * nn.MSELoss(size_average=True)(width*coord_mask, tw*coord_mask)/2.0
            loss_h = self.coord_scale * nn.MSELoss(size_average=True)(height*coord_mask, th*coord_mask)/2.0
            loss_conf = nn.MSELoss(size_average=True)(conf*conf_mask, tconf*conf_mask)/2.0
            loss_cls = self.class_scale * nn.CrossEntropyLoss(size_average=True)(cls, tcls)
            loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
            t4 = time.time()
            if False:
                print('-----------------------------------')
                print('        activation : %f' % (t1 - t0))
                print(' create pred_boxes : %f' % (t2 - t1))
                print('     build targets : %f' % (t3 - t2))
                print('       create loss : %f' % (t4 - t3))
                print('             total : %f' % (t4 - t0))
            print('%d: nGT %d, recall %d, proposals %d, loss: x %f, y %f, w %f, h %f, conf %f, cls %f, total %f' % (self.seen, nGT, nCorrect, nProposals, loss_x.item(), loss_y.item(), loss_w.item(), loss_h.item(), loss_conf.item(), loss_cls.item(), loss.item()))
            return loss
#            losses.append(loss)
#            return sum(losses).cuda()
        else:
            masked_anchors = []
            for m in self.anchor_mask:
                masked_anchors += self.anchors[m*self.anchor_step:(m+1)*self.anchor_step]
            masked_anchors = [anchor/self.stride for anchor in masked_anchors]
            boxes = utils.get_region_boxes(output.data, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask))
            return boxes

class YoloLayer2(nn.Module):
    def __init__(self, num_classes, anchors, masked_anchors, max_boxes,
                 net_width, net_height):
        super(YoloLayer2, self).__init__()

        self.num_classes = num_classes
        self.bbox_attribs = 5 + num_classes
        self.net_width = net_width
        self.net_height = net_height
        self.ignore_thresh = 0.5
        # anchors and masked anchors
        self.num_anchors = int(len(anchors)/2)
        self.anchors = anchors

        self.masked_anchor_inds = masked_anchors
        masked_anchors = []
        for i in self.masked_anchor_inds:
            masked_anchors.append(self.anchors[2*i])
            masked_anchors.append(self.anchors[2*i+1])
        self.masked_anchors = masked_anchors
        self.mask_size = int(len(masked_anchors)/2)

        self.max_boxes = max_boxes
        #self.truths = max_boxes * (4 + 1) probably dont need it
        #self.all_losses = self.batch_size * self.lwidth * self.lheight * self.lfilters # maximum number of losses to pay attention to, for a detection layer
        self.cls_loss = nn.BCELoss(size_average=False, reduce=False)
        self.l1_loss = nn.L1Loss(size_average=False, reduce=False)

    def forward(self, input, targets=None):
        # layer sizes
        batch_size = input.size(0)
        lwidth = input.size(3)
        lheight = input.size(2)
        lfilters = input.size(1)

        # reshape predictions to batch * 3 * 25 (for voc) * grid_x * grid_y
        prediction = input.view(batch_size, self.mask_size, self.bbox_attribs,
                                lheight, lwidth)
        # permute to have information on the last dimension
        prediction = prediction.permute(0,1,3,4,2).contiguous()

        x = torch.sigmoid(prediction[...,0])
        y = torch.sigmoid(prediction[...,1])
        width = prediction[...,2]
        height = prediction[...,3]
        conf = torch.sigmoid(prediction[..., 4])
        pred_cls = torch.sigmoid(prediction[..., 5:])

        if targets is not None:
            yolo_boxes = self.get_yolo_boxes(x, y, width, height, self.masked_anchors)
            targets = targets.view(batch_size, -1, 5)

            loss_objectness = torch.zeros(batch_size, self.mask_size, lheight, lwidth).cuda()
            for b in range(batch_size):
                preds = yolo_boxes[b]
                best_ious = torch.zeros(batch_size, self.mask_size, lheight, lwidth).cuda()
                for t in targets[b]:
                    if t.sum() == 0:
                        continue
                    ious = utils.bbox_ious(preds.permute(3,0,1,2), t[1:], False)
                    best_ious = torch.max(best_ious, ious)
                loss_objectness[b] = -conf[b] # we want loss to be 1 when a box is not with 100% confidence and 0

            loss_x = torch.zeros(1).cuda()
            loss_y = torch.zeros(1).cuda()
            loss_width = torch.zeros(1).cuda()
            loss_height = torch.zeros(1).cuda()
            loss_cls = torch.zeros(1).cuda()

            num_truths = 0
            for b in range(batch_size):
                for t in targets[b]:
                    if t.sum() == 0:
                        continue
                    num_truths += 1

                    gt_i = int(t[1] * lwidth)
                    gt_j = int(t[2] * lheight)

                    truth = torch.tensor([0, 0, t[3], t[4]]).cuda()
                    best_anchor = -1
                    best_iou = 0
                    for anc in range(len(self.anchors)//2):
                        anchor_box = torch.tensor([0,0,
                                                   self.anchors[2*anc]/self.net_width,
                                                   self.anchors[2*anc+1]/self.net_height]).cuda()
                        iou = utils.bbox_ious(truth, anchor_box, False)
                        if iou > best_iou:
                            best_iou = iou
                            best_anchor = anc
                    if best_anchor in self.masked_anchor_inds:
                        best_anchor_norm = best_anchor % self.mask_size
                        target_pred = yolo_boxes[b][best_anchor_norm][gt_j][gt_i]
                        iou = utils.bbox_ious(t[1:], target_pred, False)

                        tx = t[1]*lwidth - gt_i
                        ty = t[2]*lheight - gt_j
                        tw = torch.log(t[3]*self.net_width/self.anchors[2*best_anchor_norm])
                        th = torch.log(t[4]*self.net_height/self.anchors[2*best_anchor_norm+1])
                        scale = 2 * t[2] * t[3]

#                        loss_x += scale * (tx - x[b][best_anchor_norm][gt_j][gt_i])
#                        loss_y += scale * (ty - y[b][best_anchor_norm][gt_j][gt_i])
#                        loss_width += scale * (tw - width[b][best_anchor_norm][gt_j][gt_i])
#                        loss_height += scale * (th - height[b][best_anchor_norm][gt_j][gt_i])
                        loss_x += scale * self.l1_loss(x[b][best_anchor_norm][gt_j][gt_i], tx)
                        loss_y += scale * self.l1_loss(y[b][best_anchor_norm][gt_j][gt_i], ty)
                        loss_width += scale * self.l1_loss(width[b][best_anchor_norm][gt_j][gt_i], tw)
                        loss_height += scale * self.l1_loss(height[b][best_anchor_norm][gt_j][gt_i], th)

                        loss_objectness[b][best_anchor_norm][gt_j][gt_i] = 1 - conf[b][best_anchor_norm][gt_j][gt_i]

                        one_hot = torch.zeros(self.num_classes).cuda()
                        one_hot[int(t[0])] = 1.
                        for c in range(self.num_classes):
                            loss_cls += self.cls_loss(pred_cls[b][best_anchor_norm][gt_j][gt_i][c], one_hot[c])

            print("Loss x {}, loss y {}, loss w {}, loss h {}, loss_cls {}".format(loss_x/num_truths, loss_y/num_truths, loss_width/num_truths, loss_height/num_truths, loss_cls))
            return torch.sum(-loss_objectness) + loss_x/num_truths + loss_y/num_truths + loss_width/num_truths + loss_height/num_truths + loss_cls
        else:
            yolo_boxes = self.get_yolo_boxes(x, y, width, height, self.masked_anchors)

            output = torch.cat(
                    (
                            yolo_boxes.view(batch_size, -1, 4),
                            conf.view(batch_size, -1, 1),
                            pred_cls.view(batch_size, -1, self.num_classes),
                    ),
                    -1,
                    )
            return output

#            FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
#            LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
#
#            stride_w = self.net_width/lwidth
#            stride_h = self.net_height/lheight
#
#            # Calculate offsets for each grid
#            grid_x = torch.linspace(0, lwidth-1, lwidth).repeat(lwidth, 1).repeat(
#                        batch_size * self.mask_size, 1, 1).view(x.shape).type(FloatTensor)
#            grid_y = torch.linspace(0, lheight-1, lheight).repeat(lheight, 1).t().repeat(
#                        batch_size * self.mask_size, 1, 1).view(y.shape).type(FloatTensor)
#
#            scaled_anchors = FloatTensor([(a_w / stride_w, a_h / stride_h) for a_w, a_h in np.array(self.masked_anchors).reshape(-1,2)])
#
#            # Calculate anchor w, h
#            anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
#            anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
#            anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, lheight * lwidth).view(width.shape)
#            anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, lheight * lwidth).view(height.shape)
#            # Add offset and scale with anchors
#            pred_boxes = FloatTensor(prediction[..., :4].shape)
#            pred_boxes[..., 0] = x.data + grid_x
#            pred_boxes[..., 1] = y.data + grid_y
#            pred_boxes[..., 2] = torch.exp(width.data) * anchor_w
#            pred_boxes[..., 3] = torch.exp(height.data) * anchor_h
#            # Results
#            _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
#            output = torch.cat((pred_boxes.view(batch_size, -1, 4) * _scale,
#                                conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
#
#            return output.data

    def get_yolo_box(self, i, j, x, y, width, height, grid_width, grid_height, net_width, net_height, anchor_x, anchor_y):
        yolo_x = (i + x)/grid_width
        yolo_y = (j + y)/grid_height
        yolo_width = torch.exp(width) * anchor_x / net_width
        yolo_height = torch.exp(height) * anchor_y / net_height
        return torch.tensor([yolo_x, yolo_y, yolo_width, yolo_height]).cuda()

    def get_yolo_boxes(self, x, y, width, height, masked_anchors):
        assert x.shape == y.shape == width.shape == height.shape
        net_width = self.net_width
        net_height = self.net_height

        batches, num_anchors, grid_w, grid_h = x.shape
        boxes = []
        for b in range(batches):
            for a in range(num_anchors):
                for j in range(grid_h):
                    for i in range(grid_w):
                        new_x = (i + x[b][a][j][i])/grid_w
                        new_y = (j + y[b][a][j][i])/grid_h
                        new_width = torch.exp(width[b][a][j][i]) * masked_anchors[a] / net_width
                        new_height = torch.exp(height[b][a][j][i]) * masked_anchors[a+1] / net_height
                        boxes.append([new_x, new_y, new_width, new_height])
        return torch.tensor(boxes).view(batches, num_anchors, grid_w, grid_h, 4).cuda()


# -*- coding: utf-8 -*-
"""
Name: Open source Intelligent Dasymetric Mapping (IDM) script

Author: Anam Khan

Date: 5/1/19

Description: Intelligent Dasymetric Mapping (IDM) disaggregates population
counts enumerated by vector source units to the spatial resolution of a
categorical ancillary raster containing classes that are indicative of the
spatial distribution of population within the source units. This script is
an open source version of the EnviroAtlas IDM toolbox developed by
Torrin Hultgren for ArcMap: https://www.epa.gov/enviroatlas/dasymetric-toolbox.
This version follows the publication by Mennis and Hultgren (2006) with the
exception that class densities for unsampled ancillary classes are calculated
using census polygons where the population estimated for sampled/preset
ancillary classes did not exceed the census population.
"""

import os, sys, json
from osgeo import gdal, ogr
import numpy as np
import pandas as pd
import geopandas as gp
import argparse as ap

#IDM function
def dasy_map (popFeat_path, popCountField, popKeyField, ancRaster_path,
              out_dir,  popAreaMin = 1, sampleMin = 3, percent = 0.95,
              uninhab_path = False, anc_nd = 0, pop_nd = 0):
    '''
    Prepare population density rasters given population and ancillary data
    through intelligent dasymetric mapping. -popFeat_path: The path to the \
    census polygons with unique identifiers and a count of the population for \
    each polygon. - popCountField: The field in the population_features that \
    stores the polygon's populations.- popKeyField: The unique identifier \
    field for each polygon in population_features. -ancRaster_path: The path \
    to the land cover raster that is used for dasymetric population mapping. \
    -out_dir: The directory where all outputs will be saved. \
    -popAreaMin: The minimum number of raster cells in a source polygon for \
    it to be considered representative of a class (default = 1). -sampleMin: \
    The minimum number of source units to ensure a representative sample for \
    a land cover class (default = 3). -percent: The minimum percent of a \
    source polygon's area that an ancillary class must cover in order for the \
    source polygon to be considered representative of that class. Enter as a \
    decimal (default = 0.95). -uninhab_path: An optional shapefile containing \
    uninhabited areas. -anc_nd: The NoData value for the ancillary raster \
    (default = 0). -pop_nd: The NoData value for the population \
    raster (default = 0)
    '''
    #Set config.json file in the script's directory to presetTable
    if __name__ == '__main__':
        presetTable = os.path.join(sys.path[0], "config.json")
    else:
        presetTable = os.path.join(sys.path[-1], "config.json")

    print ('population_features path: {0}'.format(popFeat_path))
    print ('population_count_field: {0}'.format(popCountField))
    print ('population_key_field: {0}'.format(popKeyField))
    print ('ancillary_raster: {0}'.format(ancRaster_path))
    print ('uninhabited_file: {0}'.format(uninhab_path))
    print ('The minimum populated area of a representative unit is ' + str(popAreaMin))
    print ('The minimum sample size is ' + str(sampleMin))
    print ('The percent is ' + str(percent))
    print ('The NoData value for the population raster is ' + str(pop_nd))
    print ('The NoData value for the ancillary raster is ' + str(anc_nd))

    #Set file names for outputs
    popRaster = os.path.join(out_dir, "PopRaster.tif")
    popWorkTable = os.path.join(out_dir, "PopTable.csv")
    dasyRaster = os.path.join(out_dir, "DasyRaster.tif")
    dasyWorkTable = os.path.join(out_dir, "DasyWorkTable.csv")
    densityRaster = os.path.join(out_dir, "DensityRaster.tif")

    """
    Set driver for raster creation and read in census population features and
    ancillary raster
    """
    rast_driver = gdal.GetDriverByName('GTiff')
    ancRaster = gdal.Open(ancRaster_path)
    popFeatures = ogr.Open(popFeat_path)
    popLayer = popFeatures.GetLayer()

    '''
    Get GeoTransform from ancillary raster: rows, columns,
    coordinates of upperleft corner for north up images, and projection.
    '''
    rows = ancRaster.RasterYSize
    cols = ancRaster.RasterXSize
    ulx = ancRaster.GetGeoTransform()[0]
    uly = ancRaster.GetGeoTransform()[3]
    anc_proj = ancRaster.GetProjection()

    """
    Create population raster from census population features using the
    GeoTransform from the ancillary raster. Set pop_nd as the NoData value
    for the population raster.
    """
    print ("Creating population raster...")
    popRast = rast_driver.Create(popRaster, cols, rows, 1,
                                gdal.GDT_Float32, options=["COMPRESS=LZW"])
    popRast.SetGeoTransform((ulx, ancRaster.GetGeoTransform()[1], 0,
                             uly, 0, ancRaster.GetGeoTransform()[5]))
    popRast.SetProjection(anc_proj)
    popRast.GetRasterBand(1).SetNoDataValue(pop_nd)
    gdal.RasterizeLayer(popRast, [1], popLayer,
                        options = ["ATTRIBUTE=" + popKeyField])
    popRast = None

    """
    Burn the NoData value from the ancillary raster into the pixels that
    overlap uninhabited areas
    """
    if uninhab_path:
        uninhab_ds = ogr.Open(uninhab_path)
        uninhabLayer = uninhab_ds.GetLayer()
        uninhab_anc = os.path.join(out_dir, "uninhab_landcover.tif")
        uninhab_rast = rast_driver.CreateCopy(uninhab_anc,
                                              gdal.Open(ancRaster_path),
                                              options=['COMPRESS=LZW'])
        gdal.RasterizeLayer(uninhab_rast, [1], uninhabLayer,
                            burn_values = [anc_nd])
        uninhab_rast = None
        uninhab_ds = None
        ancRaster = gdal.Open(uninhab_anc)

    print ("Creating dasymetric units...")
    #Read ancillary raster and population raster as array
    anc_arr = ancRaster.GetRasterBand(1).ReadAsArray().astype(np.uint64)
    popRast = gdal.Open(popRaster)
    pop_arr = popRast.GetRasterBand(1).ReadAsArray().astype(np.uint64)

    """
    Convert pixel values of ancillary raster that do not overlap with census
    polygons to NoData.
    """
    anc_arr[pop_arr == pop_nd] = anc_nd

    """
    Combine the ancillary raster and the population raster using Cantor's
    pairing function to return a unique integer value for a pair(x,y)
    """
    comb_arr = 0.5 * (pop_arr + anc_arr) * (pop_arr + anc_arr + 1) + anc_arr

    """
    Write the combine array to the dasymetric raster using the GeoTransform
    from the ancillary raster
    """
    dasyRast = rast_driver.Create(dasyRaster, cols, rows, 1,
                            gdal.GDT_Float64, options=['COMPRESS=LZW'])
    dasyRast.SetGeoTransform((ulx, ancRaster.GetGeoTransform()[1], 0,
                              uly, 0, ancRaster.GetGeoTransform()[5]))
    dasyRast.SetProjection(anc_proj)
    dasyRast_b1 =dasyRast.GetRasterBand(1)
    dasyRast_b1.WriteArray(comb_arr)

    """
    Set the NoData value for the dasymetric raster band by running the same
    Cantor's pairing function on the NoData values from the population raster
    and ancillary raster.
    """
    dasy_nd = 0.5 * (pop_nd + anc_nd) * (pop_nd + anc_nd + 1) + anc_nd
    dasyRast_b1.SetNoDataValue(dasy_nd)
    dasyRast = None

    """
    Make the population DataFrame and the dasymetric DataFrame: collect the
    unique values and counts, rearrange the array via transpose, convert to
    DataFrame, and rename columns.
    """
    dasy_ar_un = np.array(np.unique(comb_arr, return_counts = True)).T
    dasy_df = pd.DataFrame(dasy_ar_un, columns = ("Value", "Count"))

    pop_ar_un = np.array(np.unique(pop_arr, return_counts = True)).T
    pop_df = pd.DataFrame(pop_ar_un, columns = ("Value", "Count"))

    '''
    Inverse of Cantor's pairing to get the polygon ID and ancillary class
    associated with each dasy unit.
    w = (sqrt(8 * dasy_df['Value'] + 1) - 1) // 2
    '''
    num = 8*dasy_df['Value'] + 1
    dasy_df['w'] = (num.pow(1./2) - 1) // 2
    dasy_df['t'] = (dasy_df['w'].pow(2) + dasy_df['w']) / 2
    dasy_df['ancID'] = dasy_df['Value'] - dasy_df['t']
    dasy_df['polyID'] = dasy_df['w'] - dasy_df['ancID']

    #get rid of unnecessary columns and NoData values
    dasy_df = dasy_df[dasy_df['Value'] != dasy_nd].drop(columns = ['w','t'])
    pop_df = pop_df[pop_df['Value'] != pop_nd]

    #Set variables for DataFrame columns
    popIDField = 'polyID'
    ancCatName = 'ancID'
    dasyAreaField = 'Count'

    #Make lists to use later
    #All ancillary categories in study area
    inAncCatList = list(np.unique(dasy_df[ancCatName]).astype(int))

    '''
    This list will be populated with ancillary categories that are not sampled
    and do not have preset class densities.
    '''
    unSampledList = []

    #Preset class densities from config.json file
    presetData = json.load(open(presetTable))

    '''
    Uninhabited classes: ancially classes where people do not live. Classes with
    a preset class density of 0
    '''
    unInhabList = [int(presetCat) for presetCat,presetVal in presetData.items()
                    if float(presetVal) == 0]

    #Ancillary classes where people can live
    InhabList = [cat for cat in inAncCatList if cat not in unInhabList]

    '''
    Join the census population counts to the dasymetric DataFrame and calculate
    population density for the polygon.
    '''
    print ("Calculating populated area...")
    #Read the source population shapefile with the population count field.
    popfeat_df = gp.read_file(popFeat_path)

    '''
    Set the polygon ID field provided by the user as an index for the
    population features DataFrame and the population DataFrame for joining and
    transfering the population count field.
    '''
    popfeat_df.index = popfeat_df[popKeyField]
    pop_df.index = pop_df["Value"]

    '''
    Join population counts from popfeat_df to the dasymetric DataFrame and the
    population DataFrame. Rename the field to "POP_COUNT" in the dasymetric
    DataFrame.
    '''
    dasy_df = dasy_df.join(popfeat_df[popCountField],
                           on = popIDField).rename(
                                   columns = {popCountField: "POP_COUNT"}
                                   )
    pop_df = pop_df.join(popfeat_df[popCountField])

    '''
    Group the dasymetric units that are associated with inhabitable classes by
    the census polygon ID and take the sum of the dasymetric area in each
    group. Rename the column as "POP_AREA".
    POP_AREA = sum(pixels) for inhabitable classes
    '''
    popAreaSum = dasy_df[
            dasy_df[ancCatName].isin(InhabList)
            ].groupby(popIDField)[dasyAreaField].sum().rename("POP_AREA")

    '''
    Transfer "POP_AREA" from popAreaSum to the dasymetric DataFrame and the
    population DataFrame.
    '''
    dasy_df["POP_AREA"] = dasy_df.join(popAreaSum, on = popIDField)["POP_AREA"]
    pop_df = pop_df.join(popAreaSum).fillna(0)

    '''
    Calculate population density for census polygons where poulated area is
    greater than 0.
    '''
    print ("Calculating population density...")
    pop_densMask = pop_df["POP_AREA"] > 0
    pop_df.loc[pop_densMask, "POP_DENS"] = pop_df.loc[
            pop_densMask, popCountField] / pop_df.loc[pop_densMask, "POP_AREA"]
    #replace NaN with 0
    pop_df = pop_df.fillna(0)

    '''
    Calculate representative population density for ancillary classes that have
    enough representative samples in the study area.
    '''
    print ("Selecting representative units...")
    #Create column for the ancillary class that a polygon is representative of
    pop_df["REP_CAT"] = 0

    '''
    For each inhabitable ancillary class, collect polygon IDs of census
    polygons that meet the user-define criteria for being representative of an
    ancillary class.
    '''
    for inAncCat in InhabList:
        repUnits_mask = (
                dasy_df["POP_AREA"] > float(popAreaMin)
                ) & (
                        dasy_df[ancCatName] == inAncCat
                        )
        repUnits = dasy_df.loc[
                repUnits_mask, [
                        dasyAreaField, popIDField, ancCatName, "POP_AREA"
                        ]
                ]
        repUnits["PERCENT"] = repUnits[dasyAreaField] / repUnits["POP_AREA"]
        repUnits = list(
                repUnits[repUnits["PERCENT"] >= float(percent)][popIDField]
                )

        if len(repUnits) >= float(sampleMin):
            pop_df.loc[pop_df['Value'].isin(repUnits), "REP_CAT"] = inAncCat
            print ("Class "
                   + str(inAncCat)
                   + " was sufficiently sampled with "
                   + str(len(repUnits))
                   + " representative source units.")

            '''
            #If ancillary category has no representative polygons and it does
            not have a preset class density, then add it to the list of
            unsampled classes.
            '''
        elif str(inAncCat) not in list(presetData):
            unSampledList.append(int(inAncCat))
            print ("Class "
                   + str(inAncCat)
                   + " was not sufficiently sampled with only "
                   + str(len(repUnits))
                   + " representative source units.")

    #Calculate statistics and make sampling summary table
    print (
            "Calculating representative population density for selected" \
            " classes..."
            )

    '''
    Create a mask for rows in the dasymetric DataFrame where REP_CAT =! 0.
    We only want to create summaries for these dasymetric rows because they are
    associated with representative polygons.
    '''
    rep_mask = pop_df["REP_CAT"] != 0

    '''
    Calculate sum of census population counts and sum of populated area for
    each sampled ancillary class.
    '''
    classDens_df = pop_df[rep_mask].groupby("REP_CAT")[
            [popCountField, 'POP_AREA']
            ].sum().rename(
            columns = {popCountField: "SUM_" + popCountField,
                       "POP_AREA": "SUM_POP_AREA"}
            )

    #Calculate sample density for sampled classes
    classDens_df["SAMPLEDENS"] = classDens_df[
            "SUM_" + popCountField
            ] / classDens_df["SUM_POP_AREA"]
    classDens_df["METHOD"] = "Sampled"
    classDens_df["CLASSDENS"] = classDens_df["SAMPLEDENS"]

    #Add preset densities to summary table
    if presetTable:
        print ("Adding preset values to the summary table...")
        for preset_cat in list(presetData):
            classDens_df.loc[int(preset_cat), "CLASSDENS"] = presetData[
                    preset_cat
                    ]
            classDens_df.loc[int(preset_cat), "METHOD"] = 'Preset'

    # For all sampled and preset classes, calculate a population estimate.
    print (
            "Calculating population estimate for sampled and preset classes..."
            )
    #Get representative population densities from class density DataFrame.
    dasy_df = dasy_df.join(classDens_df['CLASSDENS'], on = ancCatName).fillna(0)

    '''
    #Set mask for dasy_df that will limit ancillary categories to those in the
    class density DataFrame.
    '''
    popEst_mask = dasy_df[ancCatName].isin(classDens_df.index)

    '''
    POP_EST = area of the dasymetric unit
    * the representative population density of the ancillary class associated
    with the dasymetric unit
    '''
    dasy_df["POP_EST"] = 0
    dasy_df.loc[popEst_mask, "POP_EST"] = dasy_df.loc[
            popEst_mask, dasyAreaField
            ] * dasy_df.loc[
                    popEst_mask, 'CLASSDENS'
                    ]

    # Intelligent areal weighting for unsampled classes
    print ("Performing intelligent areal weighting for unsampled classes...")
    if unSampledList:
        '''
        Calculate representative population densities for unsampled ancillary
        classes using IAW
        '''
        unsampled_mask = dasy_df[ancCatName].isin(unSampledList)

        '''
        Populate remainining area of each dasymetric unit as the area of
        dasymetric units associated with unsampled classes and 0 everywhere
        else.
        '''
        dasy_df["REM_AREA"] = 0
        dasy_df.loc[unsampled_mask, "REM_AREA"] = dasy_df.loc[
                unsampled_mask, dasyAreaField
                ]

        '''
        For each polygon, sum the remaining area and sum the population that
        has already been estimated for sampled/preset classes.
        '''
        popEstSum = dasy_df.groupby(popIDField)[
                ["POP_EST", "REM_AREA"]
                ].sum()

        '''
        Join popEstSum to dasy_df to transfer the sum of population estimates
        and the sum of remaining area to the dasymetric DataFrame.
        '''
        dasy_df = dasy_df.join(popEstSum["POP_EST"], on = popIDField,
                               rsuffix = "poly")
        dasy_df = dasy_df.join(popEstSum["REM_AREA"], on = popIDField,
                               rsuffix = "poly")

        '''
        Calcualte a population difference between the census population and the
        population estimated for sampled/preset ancillary classes.
        '''
        dasy_df["POP_DIFF"] = dasy_df["POP_COUNT"] - dasy_df["POP_ESTpoly"]

        '''
        Calculate an initial population estimate for dasymetric units
        associated with unsampled ancillary classes and polygons where the
        sampled/preset population estimates did not exceed the census
        population count.
        '''
        diff_mask = (dasy_df[ancCatName].isin(unSampledList) &
                     dasy_df['REM_AREApoly'] !=0 )
        dasy_df.loc[diff_mask, "POP_EST"] = (
                dasy_df.loc[diff_mask, "POP_DIFF"].clip(0) *
                dasy_df.loc[diff_mask, "REM_AREA"] /
                dasy_df.loc[diff_mask, "REM_AREApoly"])
        '''
        Sum total initial population estimates and remaining area for
        dasymetric units used to calculate initial population estimates for
        unsampled ancillary classes.
        '''
        ancCat_sum = dasy_df[diff_mask].groupby(ancCatName)[
                ["POP_EST" , "REM_AREA"]
                ].sum()

        '''
        Calculate the representative population density for unsampled classes
        using ancCat_sum and update the class density DataFrame.
        '''
        for cat in ancCat_sum.index:
            classDens_df.loc[cat, "CLASSDENS"] = ancCat_sum.loc[cat,
                            "POP_EST"] / ancCat_sum.loc[cat,
                                           "REM_AREA"]
            classDens_df.loc[cat, "METHOD"] = "IAW"

        '''
        Add representative population densities for unsampled classes in the
        dasymetric DataFrame.
        '''
        dasy_df.loc[unsampled_mask, 'CLASSDENS'] = dasy_df.loc[
                unsampled_mask
                ].join(classDens_df.loc[
                        ancCat_sum.index, 'CLASSDENS'
                        ],
                on = ancCatName, rsuffix = "_classDens")['CLASSDENS_classDens']

        '''
        Calculate new population estimates using representative population
        densities for unsampled classes.
        POP_EST = dasymetric area * class density
        '''
        dasy_df.loc[unsampled_mask, "POP_EST"] = dasy_df.loc[unsampled_mask,
                   dasyAreaField] * dasy_df.loc[unsampled_mask,
                               'CLASSDENS']

        # End of intelligent areal weighting

    # Perform final calculations to ensure pycnophylactic integrity
    print (
            "Performing final calculations to ensure pycnophylactic" \
            " integrity..."
            )
    '''
    For each dasymetric unit, use the ratio of the estimated population to the
    total population estimated for the polygon associated with the dasymetric
    unit to redistribute the census population.
    '''

    '''
    if the sum of population densities within the source unit is equal to 0
    set the POP_EST for those to 1 (i.e., area weighting (equation 5))
    '''

    idx = (dasy_df
            .groupby(popIDField)
            .filter(
                lambda s: s['POP_EST'].sum() == 0 and
                          s['POP_COUNT'].sum() > 0
                    ).index
            )

    dasy_df.loc[idx, 'POP_EST'] = 1

    #Sum population estimates by polygon.
    popEstsum = dasy_df.groupby(popIDField)["POP_EST"].sum()

    dasy_df["TOTALFRACT"] = dasy_df["POP_EST"] / dasy_df.join(popEstsum,
           on = popIDField, rsuffix = "SUM")["POP_ESTSUM"]
    dasy_df["NEW_POP"] = dasy_df["TOTALFRACT"] * dasy_df["POP_COUNT"]
    dasy_df["NEWDENSITY"] = dasy_df["NEW_POP"] / dasy_df[dasyAreaField]

    #Replace nan with 0
    dasy_df = dasy_df.fillna(0)

    #export dasy table to .csv
    dasy_df.to_csv(dasyWorkTable, header = True)

    #export pop_df to .csv
    pop_df.fillna(0).to_csv(popWorkTable, header = True)

    #export classDens_df to sampling summary table
    classDens_df.to_csv(os.path.join(out_dir, "SamplingSummaryTable.csv"),
                        header = True)

    #Create final population density raster.
    print ("Creating population density raster...")
    #Create population density array.
    dasy_lut = dasy_df[['Value', 'NEWDENSITY']].set_index('Value')
    dasy_lut.loc[dasy_nd, 'NEWDENSITY'] = -999 #NoData value from comb_arr
    dens_df = pd.DataFrame(np.ravel(comb_arr)).join(dasy_lut, on = 0)
    dens_ar = np.array(
            dens_df['NEWDENSITY']
            ).reshape(
                    (comb_arr.shape[0], comb_arr.shape[1])
                    )

    #Write array to population density raster.
    densRast = rast_driver.Create(densityRaster, cols, rows, 1,
                                  gdal.GDT_Float32, options=['COMPRESS=LZW'])
    densRast.SetGeoTransform((ulx, ancRaster.GetGeoTransform()[1], 0,
                              uly, 0, ancRaster.GetGeoTransform()[5]))
    densRast.SetProjection(anc_proj)
    densRast_b1 =densRast.GetRasterBand(1)
    densRast_b1.WriteArray(dens_ar)
    densRast_b1.SetNoDataValue(-999)
    densRast = None

    print ("All outputs from this tool can be found in " + out_dir)

#------------------------------------------------------------------------------
#Get arguments to run dasy_pop from command line
if __name__ == '__main__':
    #create ArgumentParser
    parser = ap.ArgumentParser(description='This script accepts population \
                               and ancillary datasets for preparing \
                               population density rasters using intelligent \
                               dasymetric mapping.')

    #add arguments
    parser.add_argument('population_features', type = str,
                        help = 'The census polygons with unique identifiers \
                        and a count of the population for each polygon')
    parser.add_argument('population_count_field', type = str,
                        help = "The field in the population_features that \
                        stores the polygon's populations")
    parser.add_argument('population_key_field', type = str,
                        help = "The unique identifier field for each polygon \
                        in population_features")
    parser.add_argument('ancillary_raster', type = str,
                        help = "The land cover raster that is used for \
                        dasymetric population mapping")
    parser.add_argument('output_directory', type = str,
                        help = "The directory where all outputs from the \
                        script will be saved.")
    parser.add_argument('--uninhabited_file', type = str, nargs='?',
                        help = "An optional feature class containing \
                        uninhabited areas")
    parser.add_argument('--minimum_sampling_area', type = int,
                        nargs='?',default = 1,
                        help = "The minimum number of raster cells in a \
                        source polygon for it to be considered representative \
                        of a class - default = 1")
    parser.add_argument('--minimum_sample', type = int, nargs='?', default = 3,
                        help = "The minimum number of source units to ensure \
                        a representative sample for a land cover class \
                        - default = 3")
    parser.add_argument('--percent', type = float, nargs='?', default = 0.95,
                        help = "The minimum percent of a source polygon's \
                        area that an ancillary class must cover in order for \
                        the source polygon to be considered representative of \
                        that class. Please enter as a decimal \
                        - default = 0.95")
    parser.add_argument('--pop_nodata', type = int, nargs='?', default = 0,
                        help = "The population_features will be converted to \
                        raster and this will be the NoData value \
                        - default = 0")
    parser.add_argument('--anc_nodata', type = int, nargs='?', default = 0,
                        help = "The NoData value for the ancillary raster \
                        - default = 0")

    #get args
    args = parser.parse_args()

    #run function
    dasy_map(
            popFeat_path = args.population_features,
            popCountField = args.population_count_field,
            popKeyField = args.population_key_field,
            ancRaster_path = args.ancillary_raster,
            popAreaMin = args.minimum_sampling_area,
            sampleMin = args.minimum_sample,
            percent = args.percent,
            uninhab_path = args.uninhabited_file,
            out_dir = args.output_directory,
            anc_nd = args.pop_nodata,
            pop_nd = args.anc_nodata
            )


# -*- coding: utf-8 -*-
"""Randk-Order-Based clustering algorithm.

I combine Approximate-Rank-Order algorithm and Chinese-Whispers algorithm
 in ROCWClustering class.
"""
# Author: Soroush Moazed <soroush.moazed@gmail.com>

from __future__ import division
from __future__ import print_function

import itertools

import numpy as np
from sklearn.neighbors import NearestNeighbors

class BaseClustering:

    def __init__(self):
        pass

    def fit_predict(self, X):
        pass

    def score(self, X, y_true):
        y_pred = self.fit_predict(X)
        pairwise_precision = self._calc_pw_precision(y_pred, y_true)
        pairwise_recall = self._calc_pw_recall(y_pred, y_true)
        pairwise_f_measure = 2 * (pairwise_precision * pairwise_recall)\
            / (pairwise_precision + pairwise_recall)
        return pairwise_f_measure, pairwise_precision, pairwise_recall

    @staticmethod
    def _calc_pw_precision(y_pred, y_true):
        unique_clusters = np.unique(y_pred)
        n_pairs = 0
        n_same_class_pairs = 0
        for cluster in unique_clusters:
            sample_indices = np.where(y_pred == cluster)[0]
            combs = np.array(list(itertools.combinations(sample_indices, 2)), dtype=np.int64)
            if not np.any(combs):
                continue
            combs_classes = y_true[combs]
            same_class_pairs = np.where(combs_classes[:, 0] == combs_classes[:, 1])[0]
            n_pairs += len(combs)
            n_same_class_pairs += len(same_class_pairs)
        pw_precision = n_same_class_pairs / n_pairs
        return pw_precision

    @staticmethod
    def _calc_pw_recall(y_pred, y_true):
        unique_classes = np.unique(y_true)
        n_pairs = 0
        n_same_cluster_pairs = 0
        for clss in unique_classes:
            sample_indices = np.where(y_true == clss)[0]
            combs = np.array(list(itertools.combinations(sample_indices, 2)), dtype=np.int64)
            if not np.any(combs):
                continue
            combs_clusters = y_pred[combs]
            same_cluster_pairs = np.where(combs_clusters[:, 0] == combs_clusters[:, 1])[0]
            n_pairs += len(combs)
            n_same_cluster_pairs += len(same_cluster_pairs)
        pw_recall = n_same_cluster_pairs / n_pairs
        return pw_recall

class ROCWClustering(BaseClustering):

    """Approximated rank-order clustering implemented using Chinese Whispers algorithm.

    Using rank-order distances generate a graph, and feed this graph to ChineseWhispers
     algorithm for clustering.
    """

    def __init__(self, k, metric, n_iteration, algorithm):
        super().__init__()
        self.k = k
        self.metric = metric
        self.n_iteration = n_iteration
        self.knn_algorithm = algorithm

    def fit_predict(self, X):
        if len(X) > self.k:
            graph = ROGraph(self.k, self.metric, algorithm=self.knn_algorithm)
        else:
            graph = ROGraph(len(X), self.metric, algorithm=self.knn_algorithm)
        adjacency_mat = graph.generate_graph(X)
        clusterer = ChineseWhispersClustering(self.n_iteration)
        labels = clusterer.fit_predict(adjacency_mat)
        return labels

class ChineseWhispersClustering:

    def __init__(self, n_iteration=5):
        self.n_iteration = n_iteration
        self.adjacency_mat_ = None
        self.labels_ = None

    def fit_predict(self, adjacency_mat):

        """Fits and returns labels for samples"""

        n_nodes = adjacency_mat.shape[0]
        indices = np.arange(n_nodes)
        labels_mat = np.arange(n_nodes)
        for _ in range(self.n_iteration):
            np.random.shuffle(indices)
            for ind in indices:
                weights = adjacency_mat[ind]
                winner_label = self._find_winner_label(weights, labels_mat)
                labels_mat[ind] = winner_label
        self.adjacency_mat_ = adjacency_mat
        self.labels_ = labels_mat
        return labels_mat

    @staticmethod
    def _find_winner_label(node_weights, labels_mat):
        adjacent_nodes_indices = np.where(node_weights > 0)[0]
        adjacent_nodes_labels = labels_mat[adjacent_nodes_indices]
        unique_labels = np.unique(adjacent_nodes_labels)
        label_weights = np.zeros(len(unique_labels))
        for ind, label in enumerate(unique_labels):
            indices = np.where(adjacent_nodes_labels == label)
            weight = np.sum(node_weights[adjacent_nodes_indices[indices]])
            label_weights[ind] = weight
        winner_label = unique_labels[np.argmax(label_weights)]
        return winner_label

class ROGraph:

    def __init__(self, k, metric, algorithm):

        self.k = k
        self.metric = metric
        self.knn_algorithm = algorithm
        self.adjacency_mat_ = None

    @property
    def adjacency_mat(self):
        return self.adjacency_mat_

    def generate_graph(self, X):
        order_lists = self._get_knns(X)
        pw_distances = self._generate_normalized_pw_distances(order_lists)
        adjacency_mat = self._generate_adjacency_mat(pw_distances)
        return adjacency_mat

    def _get_knns(self, X):

        """Generates order lists and absolute distances of k-nearest-neighbors
            for each data point.
        """

        nbrs = NearestNeighbors(n_neighbors=self.k,
                                algorithm=self.knn_algorithm,
                                metric=self.metric).fit(X)
        _, order_lists = nbrs.kneighbors(X)
        return order_lists

    def _generate_normalized_pw_distances(self, order_lists):
        n_samples = len(order_lists)
        combs = itertools.combinations([i for i in range(n_samples)], 2)
        pw_distances = np.zeros((n_samples, n_samples))
        for ind1, ind2 in combs:
            order_list_1, order_list_2 = order_lists[ind1], order_lists[ind2]
            pw_dist = self._calc_pw_dist(ind1, ind2, order_list_1, order_list_2)
            pw_distances[ind1, ind2] = pw_dist
        pw_distances = pw_distances / pw_distances.max()
        pw_distances = pw_distances + pw_distances.T
        return pw_distances

    def _generate_adjacency_mat(self, pw_distances):
        adjacency_mat = self._dist2adjacency(pw_distances)
        self.adjacency_mat_ = adjacency_mat
        return adjacency_mat

    @staticmethod
    def _dist2adjacency(distances):
        mask_mat = np.zeros(distances.shape)
        mask_mat[np.where(distances > 0)] = 1
        adjacency_mat = (1 - distances) * mask_mat
        return adjacency_mat

    def _calc_pw_dist(self, ind_a, ind_b, order_list_a, order_list_b):
        pw_dist = 0.0
        if np.any(np.intersect1d(order_list_a, order_list_b)):
            order_b_in_a, order_a_in_b = self._calc_orders(ind_a, ind_b, order_list_a, order_list_b)
            d_m_ab = self._calc_dm(order_list_a, order_list_b, order_b_in_a)
            d_m_ba = self._calc_dm(order_list_b, order_list_a, order_a_in_b)
            min_of_two = min(order_a_in_b, order_b_in_a)
            pw_dist = (d_m_ab + d_m_ba) / min_of_two
        return pw_dist

    def _calc_orders(self, ind_a, ind_b, order_list_a, order_list_b):
        order_b_in_a = np.where(order_list_a == ind_b)[0]
        if not np.any(order_b_in_a):
            order_b_in_a = self.k
        else:
            order_b_in_a = order_b_in_a[0]
        order_a_in_b = np.where(order_list_b == ind_a)[0]
        if not np.any(order_a_in_b):
            order_a_in_b = self.k
        else:
            order_a_in_b = order_a_in_b[0]
        return order_b_in_a, order_a_in_b

    def _calc_dm(self, order_list_a, order_list_b, order_b_in_a):
        dist = 0
        for i in range(min(self.k, order_b_in_a)):
            sample_index = order_list_a[i]
            if np.any(order_list_b == sample_index):
                dist += 1 / self.k
            else:
                dist += 1
        return dist