import os
import sys
import pylab
import numpy
import torch
import time
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.gridspec as gridspec
import matplotlib.cm as cm
import copy

from scipy.spatial.transform import Rotation as R

def calculate_contours(pdf_vals, bin_volumes, probs=[0.68, 0.95]):
    totsum = 0.0
    flattend_pdf = pdf_vals.flatten()
    #flattend_volumes = bin_volumes.flatten()

    sorta = numpy.argsort(flattend_pdf)[::-1]

    contour_values = []

    cur_prob_index = 0

    for ind, pdf_eval in enumerate(flattend_pdf[sorta]):
        totsum += pdf_eval*bin_volumes#flattend_volumes[sorta][ind]

        if (totsum > probs[cur_prob_index]):
            contour_values.append(pdf_eval)
            cur_prob_index += 1

        if (cur_prob_index >= len(probs)):
            break

    return contour_values

def get_bounds_from_contour(cres, boundary=0.1):

    cont_min_x = 9999999.9
    cont_max_x = -9999999.9

    cont_min_y = 999999999.9
    cont_max_y = -9999999999.9

    for i in cres.allsegs[0]:

        for j in i:

            if (j[0] < cont_min_x):
                cont_min_x = j[0]
            if (j[0] > cont_max_x):
                cont_max_x = j[0]

            if (j[1] < cont_min_y):
                cont_min_y = j[1]
            if (j[1] > cont_max_y):
                cont_max_y = j[1]

    return cont_min_x, cont_max_x, cont_min_y, cont_max_y

def get_minmax_values(samples):
    mins_maxs = []

    for ind in range(samples.shape[1]):
        min_val = min(samples[:, ind])
        max_val = max(samples[:, ind])

        mins_maxs.append((min_val, max_val))

    return mins_maxs

def get_pdf_on_grid(mins_maxs, npts, model, conditional_input=None, s2_norm="standard", s2_rotate_to_true_value=False, true_values=None):

    side_vals = []

    bin_volumes = 1.0#numpy.ones([npts]*len(mins_maxs))
    glob_ind = 0
    #has_high_dim_spheres = False
    cinput = None

    sin_zen_mask=[]

    for pdf in model.pdf_defs_list:
        this_sub_dim = int(pdf[1])
        if (pdf == "s2" and s2_norm=="lambert"):
            #has_high_dim_spheres = True
            side_vals.append(numpy.linspace(-2, 2, npts))
            bin_volumes *= (side_vals[-1][1] - side_vals[-1][0])

            side_vals.append(numpy.linspace(-2, 2, npts))
            bin_volumes *= (side_vals[-1][1] - side_vals[-1][0])

            sin_zen_mask.extend([0,0])

        elif(pdf=="s2" and s2_norm=="standard"):

          sin_zen_mask.extend([1,0])

          zen_vals=numpy.linspace(mins_maxs[glob_ind][0]+1e-4, mins_maxs[glob_ind][1]-1e-4, npts)
          side_vals.append(zen_vals)
          bin_volumes*=(side_vals[-1][1] - side_vals[-1][0])

          side_vals.append(numpy.linspace(1e-4, 2*numpy.pi-1e-4, npts))
          bin_volumes *= (side_vals[-1][1] - side_vals[-1][0])

        elif(pdf=="s2"):
          raise Exception("s2_norm ", s2_norm, " unknown .")
        else:

            for ind, mm in enumerate(mins_maxs[glob_ind:glob_ind +
                                               this_sub_dim]):

                side_vals.append(numpy.linspace(mm[0], mm[1], npts))
                bin_volumes *= (side_vals[-1][1] - side_vals[-1][0])

                sin_zen_mask.append(0)

        glob_ind += this_sub_dim

    eval_positions = numpy.meshgrid(*side_vals)

    torch_positions = torch.from_numpy(
        numpy.resize(
            numpy.array(eval_positions).T,
            (npts**len(mins_maxs), len(mins_maxs))))
    eval_positions = torch_positions.clone()

    mask_inner = torch.ones(len(torch_positions)) == 1

    for ind, pdf_def in enumerate(model.pdf_defs_list):

        if (pdf_def == "s2" and s2_norm=="lambert"):

            fix_point=None

            if(s2_rotate_to_true_value and true_values is not None):
              fix_point=true_values[model.target_dim_indices[ind][0]:model.target_dim_indices[ind][1]]

            mask_inner = mask_inner & (torch.sqrt(
                (eval_positions[:, model.target_dim_indices[ind][0]:model.
                                target_dim_indices[ind][1]]**2).sum(axis=1)) <
                                       2)
            ## transform s2 subdimensions from equal-area lambert dimension to real spherical dimensiosn the model can use

            eval_positions[:, model.target_dim_indices[ind][0]:model.
                           target_dim_indices[ind]
                           [1]] = cartesian_lambert_to_spherical(
                               eval_positions[:, model.
                                              target_dim_indices[ind][0]:model.
                                              target_dim_indices[ind][1]], fix_point=fix_point)

    if (conditional_input is not None):
        cinput = conditional_input.repeat(npts**len(mins_maxs), 1)[mask_inner]

    log_res, _, _ = model(eval_positions[mask_inner], conditional_input=cinput)

    ## no conditional input and only s2 pdf .. mask bad regions
    flagged_coords=numpy.array([])
    if(conditional_input is None and model.pdf_defs_list[0]=="s2"):

      problematic_pars=model.layer_list[0][0].return_problematic_pars_between_hh_and_intrinsic(eval_positions[mask_inner], flag_pole_distance=0.02)

      if(problematic_pars.shape[0]>0):
        if(s2_norm=="lambert"):
          fix_point=None
          if(s2_rotate_to_true_value and true_values is not None):
              fix_point=true_values[model.target_dim_indices[ind][0]:model.target_dim_indices[ind][1]]
          problematic_pars=spherical_to_cartesian_lambert(problematic_pars, fix_point=fix_point)
      flagged_coords=problematic_pars.detach().numpy()

    """
    lr_mask=numpy.exp(log_res)>1e4

    print("############################# TEST")
    bad_res,_,_=model(eval_positions[mask_inner][lr_mask][:1], conditional_input=None)

    print(bad_res)

    """

    res = (-600.0)*torch.ones(len(torch_positions)).type_as(torch_positions)
    res[mask_inner] = log_res  #.exp()

    res = res.detach().numpy()
    numpy_positions=eval_positions.detach().numpy()
    if((numpy.isfinite(res)==False).sum()>0):
      print("Non-finite evaluation during PDF eval for plotting..")
      print((numpy.isfinite(res)==False).sum())
      print(numpy_positions[(numpy.isfinite(res)==False)])

      r,_,_=model(eval_positions[mask_inner][torch.isfinite(log_res)==False][:1], conditional_input=cinput)
      print(r)
      raise Exception()

    #######################

    res.resize([npts] * len(mins_maxs))

    resized_torch_positions = torch_positions.detach().numpy()
    resized_torch_positions.resize([npts] * len(mins_maxs) + [len(mins_maxs)])

    ## add in sin(theta) factors into density

    for ind, sz in enumerate(sin_zen_mask):
      if(sz==1):
        slice_mask=(None,)*ind+(slice(None,None),)+(None,)*(len(sin_zen_mask)-1-ind)

        zen_vals=numpy.sin(numpy.linspace(mins_maxs[ind][0]+1e-4, mins_maxs[ind][1]-1e-4, npts))

        ## log val, adding zenith factors where needed
        res+=numpy.log(zen_vals[slice_mask])

    return resized_torch_positions, res, bin_volumes, sin_zen_mask, flagged_coords

def rotate_coords_to(theta, phi, target, reverse=False):

  target_theta=target[0]
  target_phi=target[1]

  x=numpy.cos(target_phi)*numpy.sin(target_theta)
  y=numpy.sin(target_phi)*numpy.sin(target_theta)
  z=numpy.cos(target_theta)

  ###########

  axis=-numpy.cross(numpy.array([x,y,z]), numpy.array([0,0,1]))
  axis_len=numpy.sqrt((axis**2).sum())
  axis/=axis_len

  rot_angle=numpy.pi-target_theta
  if(reverse):
    rot_angle=-rot_angle

  axis*=rot_angle.item()

  rot_matrix = R.from_rotvec(axis)
  ###########

  x=numpy.cos(phi)*numpy.sin(theta)
  y=numpy.sin(phi)*numpy.sin(theta)
  z=numpy.cos(theta)
  vals=torch.cat([x[:,None], y[:,None],z[:,None]], dim=1)

  res=torch.from_numpy(rot_matrix.apply(vals))

  ##########

  theta=numpy.arccos(res[:,2])
  non_finite_mask=numpy.isfinite(theta)==False
  larger=non_finite_mask & (res[:,2] > 0)
  smaller=non_finite_mask & (res[:,2] < 0)

  theta[smaller]=numpy.pi
  theta[larger]=0.0

  phi=numpy.arctan2(res[:,1],res[:,0])

  #phi_smaller_mask=phi<0
  #phi[phi_smaller_mask]=phi[phi_smaller_mask]+2*numpy.pi

  return theta, phi

def cartesian_lambert_to_spherical(xl, fix_point=None):

    ## first go to spherical lambert

    r = torch.sqrt((xl**2).sum(axis=1))
    phi = torch.acos(xl[:, 0] / r)
    larger_mask = (xl[:, 1] >= 0)

    phi = larger_mask * phi + (larger_mask == 0) * (2 * numpy.pi - phi)
    theta = 2 * torch.acos(r / 2.0)

    if(fix_point is not None):

      theta, phi = rotate_coords_to(theta, phi, fix_point, reverse=True)

    ## now go to spherical real coordinates

    return torch.cat([theta[:, None], phi[:, None]], dim=1)

def spherical_to_cartesian_lambert(spherical, fix_point=None):

    #####################

    theta = spherical[:, 0]
    phi_lambert = spherical[:, 1]

    ######################
    if(fix_point is not None):
      theta, phi_lambert = rotate_coords_to(theta, phi_lambert, fix_point)

    #print(theta, phi_lambert)
    ## first go to spherical lambert
    r_lambert = 2 * torch.cos(theta / 2.0)

    x_l = r_lambert * torch.cos(phi_lambert)
    y_l = r_lambert * torch.sin(phi_lambert)

    return torch.cat([x_l[:, None], y_l[:, None]], dim=1)

def show_sample_contours(ax,
                         samples,
                         bins=50,
                         color="white",
                         contour_probs=[0.68, 0.95],
                         sin_zen_mask=[0,0]):

    ## if bins is a list, make sure samples are within bounds by artifical intervention

    new_samples = samples
    if (type(bins) == list):

        bounds_x_mask = (samples[:, 0] < bins[0][-1]) & (samples[:, 0] >
                                                         bins[0][0])
        bounds_y_mask = (samples[:, 1] < bins[1][-1]) & (samples[:, 1] >
                                                         bins[1][0])

        total_included_mask = bounds_x_mask & bounds_y_mask

        if (total_included_mask.sum() < 100):
            print("TOTAL INCLUDED SUM ", total_included_mask.sum())
            print(
                "too few SAMPLES IN PLOTTING RANGE ... fake entries to evaluate fake contour"
            )
            somex = numpy.random.uniform(bins[0][0],
                                         bins[0][-1],
                                         size=len(samples))
            somey = numpy.random.uniform(bins[1][0],
                                         bins[1][-1],
                                         size=len(samples))

            new_samples = numpy.zeros_like(samples)
            new_samples[:, 0] = somex
            new_samples[:, 1] = somey

    bin_fillings, xedges, yedges = numpy.histogram2d(new_samples[:, 0],
                                                     new_samples[:, 1],
                                                     bins=bins,
                                                     density=True)

    xvals = 0.5 * (xedges[1:] + xedges[:-1])
    yvals = 0.5 * (yedges[1:] + yedges[:-1])
    bw = (xedges[1] - xedges[0]) * (yedges[1] - yedges[0])

    bin_volumes=bw#*numpy.ones_like(bin_fillings)

    """
    for ind, m in enumerate(sin_zen_mask):
      if(m==1):
        if(ind==0):
          print("IND 0 ")
          sin_zen=numpy.sin(0.5*(xedges[1:]+xedges[:-1]))
          print(sin_zen)
          bin_volumes*=sin_zen[:,None]
        if(ind==1):
          sin_zen=numpy.sin(0.5*(yedges[1:]+yedges[:-1]))
          bin_volumes*=sin_zen[None,:]
    """

    contour_values = calculate_contours(bin_fillings, bin_volumes, probs=contour_probs)
    ## reverse
    contour_values = contour_values[::-1]

    bounds=None

    try:
      ret = ax.contour(xvals,
                       yvals,
                       bin_fillings.T,
                       levels=contour_values,
                       colors=color)

      fmt_dict = dict()

      for ind, cprob in enumerate(contour_probs[::-1]):
          if(ind<len(contour_values)):
              fmt_dict[contour_values[ind]] = "%d" % (int(cprob * 100)) + r" %"

      ax.clabel(ret,
                fontsize=9,
                inline=1,
                fmt=fmt_dict,
                levels=contour_values,
                colors=color)

      bounds = get_bounds_from_contour(ret)
    except:
      return [[-1.0,1.0], [-1.0,1.0]]

    return [[bounds[0], bounds[1]], [bounds[2], bounds[3]]]

def plot_density_with_contours(ax,
                               log_evals,
                               evalpositions,
                               bin_volumes,
                               pts_per_dim,
                               color="black",
                               contour_probs=[0.68, 0.95]):

    xvals=evalpositions[:,0,0]
    yvals=evalpositions[0,:,1]

    #xr = (evalpositions[:, :, 0].min(), evalpositions[:, :, 0].max())
    #yr = (evalpositions[:, :, 1].min(), evalpositions[:, :, 1].max())

    xr = (xvals[0], xvals[-1])
    yr = (yvals[0], yvals[-1])

    contour_values = calculate_contours(numpy.exp(log_evals), bin_volumes,
                                        probs=contour_probs)

    ## reverse

    contour_values = contour_values[::-1]

    pcol_result = ax.pcolorfast(xr, yr, numpy.exp(log_evals).T)

    #contour_x = numpy.linspace(evalpositions[:, :, 0].min(),
    #                           evalpositions[:, :, 0].max(), pts_per_dim)
    #contour_y = numpy.linspace(evalpositions[:, :, 1].min(),
    #                           evalpositions[:, :, 1].max(), pts_per_dim)

    valid_contour=True

    if(len(contour_values)>=2):
      if(contour_values[0]==contour_values[1]):
        valid_contour=False

    if(len(contour_probs) != len(contour_values)):
      valid_contour=False

    if(valid_contour):
      contour_x=xvals
      contour_y=yvals

      bins = [contour_x, contour_y]

      res = ax.contour(contour_x,
                       contour_y,
                       numpy.exp(log_evals).T,
                       levels=contour_values,
                       colors=color)

      fmt_dict = dict()

      for ind, cprob in enumerate(contour_probs[::-1]):
          fmt_dict[contour_values[ind]] = "%d" % (int(cprob * 100)) + r" %"

      ax.clabel(res,
                fontsize=9,
                inline=1,
                fmt=fmt_dict,
                levels=contour_values,
                colors=color)

    pylab.colorbar(pcol_result, ax=ax)

def get_basic_gridlines():

  n_theta=5
  n_phi=10

  gridlines=[]

  for g in numpy.linspace(0.1,numpy.pi-0.1, n_theta):
    azis=torch.linspace(0,2*numpy.pi, 100)
    zens=torch.ones_like(azis)*g
    gl=torch.cat( [zens[:,None], azis[:,None]], dim=1)
    gridlines.append(gl)

  for a in numpy.linspace(0,2*numpy.pi-2*numpy.pi/n_phi, n_phi):
    zens=torch.linspace(0,numpy.pi,100)
    azis=torch.ones_like(zens)*a
    gl=torch.cat( [zens[:,None], azis[:,None]], dim=1)
    gridlines.append(gl)

  return gridlines

def plot_joint_pdf(pdf,
                   fig,
                   gridspec,
                   samples,
                   subgridspec=None,
                   conditional_input=None,
                   bounds=None,
                   multiplot=False,
                   total_pdf_eval_pts=10000,
                   true_values=None,
                   plot_only_contours=False,
                   contour_probs=[0.68, 0.95],
                   contour_color="white",
                   autoscale=True,
                   skip_plotting_density=False,
                   hide_labels=False,
                   s2_norm="standard",
                   colormap=cm.rainbow,
                   s2_rotate_to_true_value=True,
                   s2_show_gridlines=True,
                   skip_plotting_samples=False):

    plot_density = False
    dim = len(samples[0])
    if (pdf.total_target_dim == 1 and dim == pdf.total_target_dim):
        plot_density = True

    if (pdf.total_target_dim == 2 and dim == pdf.total_target_dim):
        plot_density = True

    if (conditional_input is not None):
        if (conditional_input.shape[0] > 1):
            plot_density = False

    if (skip_plotting_density):
        plot_density = False

    mms = get_minmax_values(samples)

    if (bounds is not None):
        mms = bounds

    ## true positions are typically labels
    plotted_true_values=None
    if(true_values is not None):
      plotted_true_values=copy.deepcopy(true_values)

    ## if bounds contain torch .. change to float

    pure_float_mms = []
    for b in mms:
        new_b = b
        if (type(b[0]) == torch.Tensor):
            new_b = [
                float(b[0].detach().numpy()),
                float(b[1].detach().numpy())
            ]

        pure_float_mms.append(new_b)

    totalpts = total_pdf_eval_pts
    pts_per_dim = int(totalpts**(1.0 / float(dim)))

    samples = samples.detach().clone()

    gridline_dict=None
    if(s2_show_gridlines and "s2" in pdf.pdf_defs_list):
      gridline_dict=dict()
      for ind, pdf_type in enumerate(pdf.pdf_defs_list):
        if(pdf_type=="s2"):
          gridline_dict[(pdf.target_dim_indices[ind][0], pdf.target_dim_indices[ind][1])]=get_basic_gridlines()

    ## transform samples to lambert space if necessary
    for ind, pdf_type in enumerate(pdf.pdf_defs_list):
        if (pdf_type == "s2" and s2_norm=="lambert"):
            ## transform to labmert space

            ## calculate fix point if rotation for visualtizion is desired
            fix_point=None

            if(s2_rotate_to_true_value and true_values is not None):

              fix_point=true_values[pdf.target_dim_indices[ind][0]:pdf.target_dim_indices[ind][1]]

            ## transform samples to lambert space
            samples[:,
                    pdf.target_dim_indices[ind][0]:pdf.target_dim_indices[ind]
                    [1]] = spherical_to_cartesian_lambert(
                        samples[:, pdf.target_dim_indices[ind][0]:pdf.
                                target_dim_indices[ind][1]], fix_point=fix_point)

            ## transform true value to lambert space
            if(plotted_true_values is not None):

                res=spherical_to_cartesian_lambert(true_values[pdf.target_dim_indices[ind][0]:pdf.
                target_dim_indices[ind][1]].unsqueeze(0), fix_point=fix_point)

                plotted_true_values[pdf.target_dim_indices[ind][0]:pdf.
                target_dim_indices[ind][1]]=res.squeeze(0)

            ## transform gridlines to lambert space
            if(s2_show_gridlines):
              tup=(pdf.target_dim_indices[ind][0],pdf.target_dim_indices[ind][1])

              new_list=[]

              for l in gridline_dict[tup]:
                new_list.append(spherical_to_cartesian_lambert(l, fix_point=fix_point))

              gridline_dict[tup]=new_list

    samples = samples.numpy()

    pdf_conditional_input = conditional_input

    if (pdf_conditional_input is not None):
        pdf_conditional_input = pdf_conditional_input[0:1]

    evalpositions, log_evals, bin_volumes, sin_zen_mask, unreliable_spherical_regions= get_pdf_on_grid(
        pure_float_mms,
        pts_per_dim,
        pdf,
        conditional_input=pdf_conditional_input,
        s2_norm=s2_norm,
        s2_rotate_to_true_value=s2_rotate_to_true_value,
        true_values=true_values)

    total_pdf_integral=numpy.exp(log_evals).sum()*bin_volumes

    if (dim == 1):
        ax = fig.add_subplot(gridspec)
        ax.hist(samples[:, 0], bins=50, density=True)

        if (plot_density):

            ax.plot(evalpositions[:, 0], numpy.exp(log_evals), color="k")

        if (true_values is not None):
            ax.axvline(true_values[0], color="red", lw=2.0)

        if (hide_labels):
            ax.set_yticklabels([])
            ax.set_xticklabels([])

    elif (dim == 2 and multiplot == False):

        if (subgridspec is None):
            subgridspec = gridspec.subgridspec(1, 1)

        ax = fig.add_subplot(subgridspec[0, 0])

        hist_bounds = 50
        #if(bounds is not None):
        #    hist_bounds=[numpy.linspace(bounds[0][0], bounds[0][1], 50),numpy.linspace(bounds[1][0], bounds[1][1], 50) ]

        ## plot the density and contours from density
        if (plot_density):

            if (contour_probs != [] and skip_plotting_samples==False):

              two_d_bounds_for_better_density = show_sample_contours(ax,
                                                samples,
                                                bins=50,
                                                color=contour_color,
                                                contour_probs=contour_probs,
                                                sin_zen_mask=sin_zen_mask)

              x_width=two_d_bounds_for_better_density[0][1]-two_d_bounds_for_better_density[0][0]
              y_width=two_d_bounds_for_better_density[1][1]-two_d_bounds_for_better_density[1][0]

              extra_x=x_width*0.2
              extra_y=y_width*0.2

              two_d_bounds_for_better_density[0][0]-=extra_x
              two_d_bounds_for_better_density[0][1]+=extra_x

              two_d_bounds_for_better_density[1][0]-=extra_y
              two_d_bounds_for_better_density[1][1]+=extra_y

              evalpositions_2d, log_evals_2d, bin_volumes_2d, _, _= get_pdf_on_grid(
                two_d_bounds_for_better_density,
                pts_per_dim,
                pdf,
                conditional_input=pdf_conditional_input,
                s2_norm=s2_norm,
                s2_rotate_to_true_value=s2_rotate_to_true_value,
                true_values=true_values)

              plot_density_with_contours(ax, log_evals_2d, evalpositions_2d,
                                         bin_volumes_2d, pts_per_dim)

            else:

              plot_density_with_contours(ax, log_evals, evalpositions,
                                         bin_volumes, pts_per_dim)

        ## plot a histogram density from samples

        if ( (plot_only_contours == False) and (plot_density == False) and (skip_plotting_samples==False)):

            ax.hist2d(samples[:, 0],
                      samples[:, 1],
                      bins=hist_bounds,
                      density=True)

        ## plot contours from samples
        new_bounds = None
        if (contour_probs != [] and skip_plotting_samples==False):

            new_bounds = show_sample_contours(ax,
                                              samples,
                                              bins=hist_bounds,
                                              color=contour_color,
                                              contour_probs=contour_probs,
                                              sin_zen_mask=sin_zen_mask)

        if (bounds is not None):
            new_bounds = bounds

        ## mark poles
        if(len(unreliable_spherical_regions)>0):

          ax.plot(unreliable_spherical_regions[:,0], unreliable_spherical_regions[:,1], color="orange", marker="x", lw=0.0)

        ## plot true values
        if (plotted_true_values is not None):

            ax.plot([plotted_true_values[0]], [plotted_true_values[1]],
                    color="red",
                    marker="o",
                    ms=3.0)

        ## plot gridlines if desired
        if(s2_show_gridlines and gridline_dict is not None):

          for gl in gridline_dict[(0,2)]:
            np_gl=gl.numpy()

            ax.plot(np_gl.T[0], np_gl.T[1], color="gray", alpha=0.5)

        ## adjust axis bounds

        if (new_bounds is not None):
            ax.set_xlim(new_bounds[0][0], new_bounds[0][1])
            ax.set_ylim(new_bounds[1][0], new_bounds[1][1])

        ### overwrite any bounds for spherical
        if(pdf.pdf_defs_list[0]=="s2"):
          if(s2_norm=="standard"):
            ax.set_xlim(0, numpy.pi)
            ax.set_ylim(0, 2*numpy.pi)
          else:
            ax.set_xlim(-2.0,2.0)
            ax.set_ylim(-2.0,2.0)

        if (hide_labels):
            ax.set_yticklabels([])
            ax.set_xticklabels([])

    else:

        if (subgridspec is None):
            subgridspec = gridspec.subgridspec(dim, dim)

        for ind1 in range(dim):
            for ind2 in range(dim):

                if (ind2 < ind1):

                    ax = fig.add_subplot(subgridspec[ind1, ind2])

                    ## make sure background looks similar to histogram empty bins
                    ax.set_facecolor(colormap(0.0))

                    hist_bounds = 50
                    #if(bounds is not None):
                    #    hist_bounds=[numpy.linspace(bounds[ind2][0], bounds[ind2][1], 50),numpy.linspace(bounds[ind1][0], bounds[ind1][1], 50) ]

                    if (plot_only_contours == False):
                        ax.hist2d(samples[:, ind2],
                                  samples[:, ind1],
                                  bins=hist_bounds,
                                  density=True,
                                  cmap=colormap)

                    if (true_values is not None):
                        ax.plot([true_values[ind2]], [true_values[ind1]],
                                color="red",
                                marker="o",
                                ms=3.0)

                    new_samples = numpy.concatenate(
                        [samples[:, ind2:ind2 + 1], samples[:, ind1:ind1 + 1]],
                        axis=1)

                    new_bounds = None
                    if (contour_probs != []):
                        new_bounds = show_sample_contours(
                            ax,
                            new_samples,
                            bins=hist_bounds,
                            color=contour_color,
                            contour_probs=contour_probs)

                    if (bounds is not None):
                        new_bounds = [bounds[ind2], bounds[ind1]]

                    if (autoscale and new_bounds is not None):

                        ax.set_xlim(new_bounds[0][0], new_bounds[0][1])
                        ax.set_ylim(new_bounds[1][0], new_bounds[1][1])

                    if (hide_labels):
                        ax.set_yticklabels([])
                        ax.set_xticklabels([])

                elif (ind2 == ind1):

                    ax = fig.add_subplot(subgridspec[ind1, ind2])

                    hist_bounds = 50
                    if (bounds is not None):
                        hist_bounds = numpy.linspace(bounds[ind2][0],
                                                     bounds[ind2][1], 50)
                    ax.hist(samples[:, ind1], bins=hist_bounds, density=True)

                    if (true_values is not None):
                        ax.axvline(true_values[ind1], color="red", lw=2.0)

                    if (autoscale):
                        if (bounds is not None):
                            ax.set_xlim(bounds[ind2][0], bounds[ind2][1])
                    if (hide_labels):
                        ax.set_yticklabels([])
                        ax.set_xticklabels([])

    return subgridspec

def visualize_pdf(pdf,
                  fig,
                  gridspec=None,
                  subgridspec=None,
                  conditional_input=None,
                  nsamples=10000,
                  total_pdf_eval_pts=10000,
                  bounds=None,
                  true_values=None,
                  plot_only_contours=False,
                  contour_probs=[0.68, 0.95],
                  contour_color="white",
                  autoscale=True,
                  seed=None,
                  skip_plotting_density=False,
                  hide_labels=False,
                  s2_norm="standard",
                  colormap=cm.rainbow,
                  s2_rotate_to_true_value=True,
                  s2_show_gridlines=True,
                  skip_plotting_samples=False):

    with torch.no_grad():
      sample_conditional_input = conditional_input
      if (conditional_input is not None):
          if (len(conditional_input.shape) == 1):
              sample_conditional_input = sample_conditional_input.unsqueeze(
                  0)

          if (sample_conditional_input.shape[0] == 1):
              sample_conditional_input = sample_conditional_input.repeat(
                  nsamples, 1)

      if (gridspec is None):
          gridspec = fig.add_gridspec(1, 1)[0, 0]

      samples, samples_base, evals, evals_base = pdf.sample(
          samplesize=nsamples,
          conditional_input=sample_conditional_input,
          seed=seed)

      higher_dim_spheres = False

      new_subgridspec = plot_joint_pdf(
          pdf,
          fig,
          gridspec,
          samples,
          subgridspec=subgridspec,
          conditional_input=conditional_input,
          bounds=bounds,
          multiplot=False,
          total_pdf_eval_pts=total_pdf_eval_pts,
          true_values=true_values,
          plot_only_contours=plot_only_contours,
          contour_probs=contour_probs,
          contour_color=contour_color,
          autoscale=autoscale,
          skip_plotting_density=skip_plotting_density,
          hide_labels=hide_labels,
          s2_norm=s2_norm,
          colormap=colormap,
          s2_rotate_to_true_value=s2_rotate_to_true_value,
          s2_show_gridlines=s2_show_gridlines,
          skip_plotting_samples=skip_plotting_samples)

    return samples, new_subgridspec


<filename>metrics/wasserstein.py<gh_stars>1-10
import numpy as np
from scipy.linalg import sqrtm

def wasserstein_2(P, Q):
    """

    :param P: (m, d) matrix. Reference data
    :param Q: (n, d) matrix. Distribution we are evaluating
    :return: float
    """

    P_mean, P_cov = np.mean(P, axis=0), np.cov(P, rowvar=False)
    Q_mean, Q_cov = np.mean(Q, axis=0), np.cov(Q, rowvar=False)

    QP_diff = Q_mean - P_mean
    return np.inner(QP_diff, QP_diff) + np.trace(Q_cov + P_cov - 2 * sqrtm(Q_cov @ P_cov))


<reponame>sharechanxd/Risk-Index-Covid19
dict_iso={'afghanistan': 'Afghanistan',
 'albania': 'Albania',
 'algeria': 'Algeria',
 'andorra': 'Andorra',
 'angola': 'Angola',
 'antigua-and-barbuda': 'Antigua and Barbuda',
 'argentina': 'Argentina',
 'armenia': 'Armenia',
 'aruba': 'Aruba',
 'australia': 'Australia',
 'austria': 'Austria',
 'azerbaijan': 'Azerbaijan',
 'bahamas': 'Bahamas',
 'bahrain': 'Bahrain',
 'bangladesh': 'Bangladesh',
 'Barbados': 'Barbados',
 'belarus': 'Belarus',
 'belgium': 'Belgium',
 'belize': 'Belize',
 'benin': 'Benin',
 'bermuda': 'Bermuda',
 'bhutan': 'Bhutan',
 'bolivia': 'Bolivia, Plurinational State of',
 'bosnia-and-herzegovina': 'Bosnia and Herzegovina',
 'botswana': 'Botswana',
 'brazil': 'Brazil',
 'bulgaria': 'Bulgaria',
 'burkina-faso': 'Burkina Faso',
 'burundi': 'Burundi',
 'cabo-verde': 'Cape Verde',
 'cambodia': 'Cambodia',
 'cameroon': 'Cameroon',
 'canada': 'Canada',
 'cayman-islands': 'Cayman Islands',
 'central-african-republic': 'Central African Republic',
 'chad': 'Chad',
 'chile': 'Chile',
 'china': 'China',
 'china-hong-kong-sar': 'Hong Kong,China',
 'china-macao-sar': 'Macao,China',
 'colombia': 'Colombia',
 'comoros': 'Comoros',
 'congo': 'Congo',
 'costa-rica': 'Costa Rica',
 'cote-d-ivoire': "Côte d'Ivoire",
 'croatia': 'Croatia',
 'cuba': 'Cuba',
 'cyprus': 'Cyprus',
 'czech-republic': 'Czech Republic',
 'democratic-republic-of-the-congo': 'Congo, the Democratic Republic of the',
 'denmark': 'Denmark',
 'djibouti': 'Djibouti',
 'dominican-republic': 'Dominican Republic',
 'ecuador': 'Ecuador',
 'egypt': 'Egypt',
 'el-salvador': 'El Salvador',
 'equatorial-guinea': 'Equatorial Guinea',
 'eritrea': 'Eritrea',
 'estonia': 'Estonia',
 'ethiopia': 'Ethiopia',
 'faeroe-islands': 'Faroe Islands',
 'fiji': 'Fiji',
 'finland': 'Finland',
 'france': 'France',
 'french-guiana': 'French Guiana',
 'french-polynesia': 'French Polynesia',
 'gabon': 'Gabon',
 'gambia': 'Gambia',
 'georgia': 'Georgia',
 'germany': 'Germany',
 'ghana': 'Ghana',
 'gibraltar': 'Gibraltar',
 'greece': 'Greece',
 'grenada': 'Grenada',
 'guadeloupe': 'Guadeloupe',
 'guatemala': 'Guatemala',
 'guinea': 'Guinea',
 'guinea-bissau': 'Guinea-Bissau',
 'guyana': 'Guyana',
 'haiti': 'Haiti',
 'honduras': 'Honduras',
 'hungary': 'Hungary',
 'iceland': 'Iceland',
 'india': 'India',
 'indonesia': 'Indonesia',
 'iran': 'Iran, Islamic Republic of',
 'iraq': 'Iraq',
 'ireland': 'Ireland',
 'israel': 'Israel',
 'italy': 'Italy',
 'jamaica': 'Jamaica',
 'japan': 'Japan',
 'jordan': 'Jordan',
 'kazakhstan': 'Kazakhstan',
 'kenya': 'Kenya',
 'kuwait': 'Kuwait',
 'kyrgyzstan': 'Kyrgyzstan',
 'latvia': 'Latvia',
 'lebanon': 'Lebanon',
 'lesotho': 'Lesotho',
 'liberia': 'Liberia',
 'libya': 'Libya',
 'liechtenstein': 'Liechtenstein',
 'lithuania': 'Lithuania',
 'luxembourg': 'Luxembourg',
 'macedonia': 'North Macedonia',
 'madagascar': 'Madagascar',
 'malawi': 'Malawi',
 'malaysia': 'Malaysia',
 'maldives': 'Maldives',
 'mali': 'Mali',
 'malta': 'Malta',
 'martinique': 'Martinique',
 'mauritania': 'Mauritania',
 'mauritius': 'Mauritius',
 'mayotte': 'Mayotte',
 'mexico': 'Mexico',
 'moldova': 'Moldova, Republic of',
 'monaco': 'Monaco',
 'mongolia': 'Mongolia',
 'montenegro': 'Montenegro',
 'morocco': 'Morocco',
 'mozambique': 'Mozambique',
 'myanmar': 'Myanmar',
 'namibia': 'Namibia',
 'nepal': 'Nepal',
 'netherlands': 'Netherlands',
 'new-zealand': 'New Zealand',
 'nicaragua': 'Nicaragua',
 'niger': 'Niger',
 'nigeria': 'Nigeria',
 'norway': 'Norway',
 'oman': 'Oman',
 'pakistan': 'Pakistan',
 'panama': 'Panama',
 'papua-new-guinea': 'Papua New Guinea',
 'paraguay': 'Paraguay',
 'peru': 'Peru',
 'philippines': 'Philippines',
 'poland': 'Poland',
 'portugal': 'Portugal',
 'qatar': 'Qatar',
 'reunion': 'Réunion',
 'romania': 'Romania',
 'russia': 'Russia',
 'rwanda': 'Rwanda',
 'saint-kitts-and-nevis': 'Saint Kitts and Nevis',
 'saint-lucia': 'Saint Lucia',
 'sao-tome-and-principe': 'Sao Tome and Principe',
 'saudi-arabia': 'Saudi Arabia',
 'senegal': 'Senegal',
 'serbia': 'Serbia',
 'seychelles': 'Seychelles',
 'sierra-leone': 'Sierra Leone',
 'singapore': 'Singapore',
 'slovakia': 'Slovakia',
 'slovenia': 'Slovenia',
 'somalia': 'Somalia',
 'south-africa': 'South Africa',
 'south-korea': 'South Korea',
 'spain': 'Spain',
 'sri-lanka': 'Sri Lanka',
 'state-of-palestine': 'Palestinian Territory, Occupied',
 'sudan': 'Sudan',
 'suriname': 'Suriname',
 'swaziland': 'Swaziland',
 'sweden': 'Sweden',
 'switzerland': 'Switzerland',
 'syria': 'Syrian Arab Republic',
 'taiwan': 'Taiwan,China',
 'tajikistan': 'Tajikistan',
 'tanzania': 'Tanzania, United Republic of',
 'thailand': 'Thailand',
 'togo': 'Togo',
 'trinidad-and-tobago': 'Trinidad and Tobago',
 'tunisia': 'Tunisia',
 'turkey': 'Turkey',
 'turks-and-caicos-islands': 'Turks and Caicos Islands',
 'uganda': 'Uganda',
 'uk': 'United Kingdom',
 'ukraine': 'Ukraine',
 'united-arab-emirates': 'United Arab Emirates',
 'uruguay': 'Uruguay',
 'us': 'United States',
 'uzbekistan': 'Uzbekistan',
 'venezuela': 'Venezuela, Bolivarian Republic of',
 'viet-nam': 'Viet Nam',
 'western-sahara': 'Western Sahara',
 'yemen': 'Yemen',
 'zambia': 'Zambia',
 'zimbabwe': 'Zimbabwe',
 'faeroe-islands':'Faroe Islands',
'saint-vincent-and-the-grenadines':'Saint Vincent & the Grenadines',
'timor-leste':'Timor-Leste',
'grenada':'Grenada',
'new-caledonia':'New Caledonia',
'laos':'Lao People\'s Democratic Republic',
'dominica':'Dominica',
'falkland-islands-malvinas':'Falkland Islands',
'greenland':'Greenland',
'holy-see':'Holy See (Vatican City State)',
'anguilla':'Anguilla',
'south-sudan':'South Sudan'
}

cate={'china':'east asia',
 'us':'north america',
 'brazil':'south america',
 'russia':'eastern europe',
 'india':'south asia',
 'uk':'western europe',
 'spain':'western europe',
 'peru':'south america',
 'chile':'south america',
 'italy':'western europe',
 'iran':'west asia',
 'mexico':'central america and mexico',
 'pakistan':'west asia',
 'turkey':'west asia',
 'germany':'western europe',
 'saudi-arabia':'west asia',
 'france':'western europe',
 'south-africa':'southern africa',
 'bangladesh':'south asia',
 'canada':'north america',
 'qatar':'west asia',
'democratic-republic-of-the-congo':'central africa',
 'colombia':'south america',
 'egypt':'south-east mediterranean',
 'sweden':'western europe',
 'belarus':'eastern europe',
 'belgium':'western europe',
 'argentina':'south america',
 'ecuador':'south america',
 'indonesia':'southeast asia',
 'netherlands':'western europe',
 'united-arab-emirates':'west asia',
 'iraq':'west asia',
 'kuwait':'west asia',
 'singapore':'southeast asia',
 'ukraine':'eastern europe',
 'portugal':'western europe',
 'oman':'west asia',
 'philippines':'southeast asia',
 'poland':'eastern europe',
 'panama':'central america and mexico',
 'switzerland':'western europe',
 'dominican-republic':'caribbean',
 'afghanistan':'west asia',
 'bolivia':'south america',
 'romania':'eastern europe',
 'bahrain':'west asia',
 'ireland':'western europe',
 'armenia':'eastern europe',
 'nigeria':'west africa',
 'israel':'south-east mediterranean',
 'kazakhstan':'central asia',
 'japan':'east asia',
 'austria':'western europe',
 'honduras':'central america and mexico',
'sao-tome-and-principe':'southeast asia',
'central-african-republic':'central africa',
'gabon':'central africa',
 'ghana':'west africa',
 'azerbaijan':'central asia',
 'guatemala':'central america and mexico',
 'moldova':'eastern europe',
 'serbia':'eastern europe',
 'algeria':'south-east mediterranean',
 'nepal':'south asia',
 'south-korea':'east asia',
 'denmark':'western europe',
 'cameroon':'central africa',
 'morocco':'south-east mediterranean',
 'czech-republic':'eastern europe',
 'sudan':'east africa',
 'cote-d-ivoire':'west africa',
 'norway':'western europe',
 'malaysia':'southeast asia',
 'uzbekistan':'central asia',
 'australia':'pacific region',
 'finland':'western europe',
 'saint-martin':'caribbean',
 'senegal':'west africa',
 'macedonia':'eastern europe',
 'kenya':'east africa',
 'el-salvador':'central america and mexico',
 'guyana':'caribbean',
 'tajikistan':'central asia',
 'ethiopia':'east africa',
 'guinea':'west africa',
 'venezuela':'south america',
 'jamaica':'caribbean',
 'kyrgyzstan':'central asia',
 'bulgaria':'eastern europe',
 'djibouti':'east africa',
 'luxembourg':'western europe',
 'mauritania':'west africa',
 'hungary':'eastern europe',
 'bosnia-and-herzegovina':'eastern europe',
 'french-guiana':'south america',
 'grenada':'caribbean',
 'greece':'western europe',
 'thailand':'southeast asia',
 'costa-rica':'central america and mexico',
 'suriname':'caribbean',
 'somalia':'east africa',
 'croatia':'eastern europe',
 'mayotte':'east africa',
 'albania':'eastern europe',
 'cuba':'caribbean',
 'maldives':'south asia',
 'nicaragua':'central america and mexico',
 'equatorial-guinea':'central africa',
 'mali':'west africa',
 'paraguay':'south america',
 'madagascar':'indian ocean islands',
 'sri-lanka':'south asia',
 'haiti':'caribbean',
 'state-of-palestine':'missing',
 'south-sudan':'east africa',
 'estonia':'eastern europe',
 'iceland':'western europe',
 'lithuania':'eastern europe',
 'lebanon':'south-east mediterranean',
 'slovakia':'eastern europe',
 'guinea-bissau':'west africa',
 'slovenia':'eastern europe',
 'zambia':'southern africa',
 'new-zealand':'pacific region',
 'sierra-leone':'west africa',
 'china-hong-kong-sar':'east asia',
 'tunisia':'south-east mediterranean',
 'cabo-verde':'west africa',
 'benin':'west africa',
 'malawi':'southern africa',
 'jordan':'south-east mediterranean',
 'yemen':'west asia',
 'latvia':'eastern europe',
 'niger':'west africa',
 'cyprus':'south-east mediterranean',
 'burkina-faso':'west africa',
 'uruguay':'south america',
 'georgia':'eastern europe',
 'rwanda':'east africa',
 'chad':'west africa',
 'mozambique':'southern africa',
 'uganda':'east africa',
 'andorra':'western europe',
 'swaziland':'southern africa',
 'liberia':'west africa',
 'libya':'south-east mediterranean',
 'malta':'south-east mediterranean',
 'togo':'west africa',
 'channel-islands':'western europe',
 'zimbabwe':'southern africa',
 'reunion':'indian ocean islands',
 'tanzania':'southern africa',
 'montenegro':'eastern europe',
 'taiwan':'east asia',
 'viet-nam':'southeast asia',
 'mauritius':'west africa',
 'myanmar':'southeast asia',
 'comoros':'indian ocean islands',
 'angola':'southern africa',
 'syria':'south-east mediterranean',
 'martinique':'eastern europe',
 'mongolia':'east asia',
 'cayman-islands':'north america',
 'eritrea':'east africa',
 'namibia':'southern africa',
 'guadeloupe':'caribbean',
 'gibraltar':'north africa',
 'burundi':'east africa',
 'bermuda':'north america',
 'cambodia':'southeast asia',
 'bahamas':'caribbean',
 'monaco':'eastern europe',
 'botswana':'southern africa',
 'bhutan':'south asia',
 'seychelles':'indian ocean islands',
 'antigua-and-barbuda':'caribbean',
 'french-polynesia':'pacific region',
 'china-macao-sar':'east asia',
 'gambia':'west africa',
 'turks-and-caicos-islands':'southern africa',
 'lesotho':'southern africa',
 'belize':'caribbean',
 'curacao':'north america',
 'papua-new-guinea':'pacific region',
 'western-sahara':'west africa',
 'fiji':'pacific region',
 'saint-kitts-and-nevis':'caribbean',
 'saint-lucia':'caribbean',
 'congo':'west africa',
 'trinidad-and-tobago':'caribbean',
 'faeroe-islands':'western europe',
 'Barbados':'caribbean',
 'liechtenstein':'western europe',
 'aruba':'western europe',
 'faeroe-islands':'western europe',
'saint-vincent-and-the-grenadines':'caribbean',
'timor-leste':'pacific region',
'grenada':'caribbean',
'new-caledonia':'pacific region',
'laos':'southeast asia',
'dominica':'caribbean',
'falkland-islands-malvinas':'south america',
'greenland':'north america',
'holy-see':'western europe',
'anguilla':'caribbean',
 }
from tqdm import tqdm
import numpy as np
import pandas as pd
from bs4 import BeautifulSoup
import requests
import json
import time
import random
import html5lib
import re
import scipy.stats as st
from pandas.core.frame import DataFrame
import copy
import math
import os
import datetime
#import datetime
x0=datetime.date.today()
x1=datetime.date.today()-datetime.timedelta(days=1)
x2=datetime.date.today()-datetime.timedelta(days=2)
# run_time
ts=[]
ts.append(x0.__format__('%Y%m%d'))
ts.append(x1.__format__('%Y%m%d'))
ts.append(x2.__format__('%Y%m%d'))
print(ts)
headers = { 'Connection': 'close',}
# proxies={'http':'http://127.0.0.1:10080','https':'http://127.0.0.1:10080'}
url='https://www.worldometers.info/coronavirus/#countries'
# url='https://www.worldometers.info/coronavirus/country/us/'
a=requests.get(url,headers=headers)
soup = BeautifulSoup(a.content,'html5lib')
x=soup.body.find_all('tr', attrs={'style': ['','background-color:#F0F0F0','background-color:#EAF7D5']})
# 190 210
def find_start_yesterday(i,j):
    for start in range(i,j):
        one=x[start]
        two=x[start+1]
        l1=one.find_all('a',attrs={'class':'mt_a'})
        l2=two.find_all('a',attrs={'class':'mt_a'})
        if l1==[] or l2==[]:
            continue
        s1=str(l1[0])
        s2=str(l2[0])
        coun1=s1.split('/')
        coun2=s2.split('/')
        if coun1[1]=='china' and coun2[1]=='us':
            return start
#385 410
# def find_end_yesterday(i,j):
#     for end in range(i,j):
#         final_pre=x[end-1]
#         final=x[end]
#         l1=final_pre.find_all('a',attrs={'class':'mt_a'})
#         l2=final.find_all('a',attrs={'class':'mt_a'})
#         if l1==[] or l2==[]:
#             continue
#         s1=str(l1[0])
#         s2=str(l2[0])
#         coun1=s1.split('/')
#         coun2=s2.split('/')
#         if (coun1[1]=='anguilla' and coun2[1]=='saint-pierre-and-miquelon') or (coun2[1]=='anguilla' and coun1[1]=='saint-pierre-and-miquelon'):
#             return end+1

def find_end_yesterday(i,j):
    for end in range(i,j):
        # final_pre=x[end-1]
        final=x[end]
        # l1=final_pre.find_all('a',attrs={'class':'mt_a'})
        l2=final.find_all('a',attrs={'class':'mt_a'})
        if l2==[]:
            continue
        # s1=str(l1[0])
        s2=str(l2[0])
        # coun1=s1.split('/')
        coun2=s2.split('/')
        if coun2[1]=='anguilla':
            return end+1

end=find_end_yesterday(400,440)
end2=find_end_yesterday(630,700)
start=find_start_yesterday(190,250)
start2=find_start_yesterday(440,470)
print('start:{}\tend:{}\tstart2:{}\tend2:{}'.format(start,end,start2,end2))

col_name=['0','#','Country,Other','TotalCases',
 'NewCases','TotalDeaths','NewDeaths','TotalRecovered',
 'NewRecovered','ActiveCases','Serious,Critical','Tot Cases/1M pop',
 'Deaths/1M pop','TotalTests','Tests/1M pop','Population',
 'Continent','17',' 1 Caseevery X', 'ppl1 Deathevery',' X ppl1 Testevery ','X ppl','22',
 'Cases Per 100K Population','Tests Per 100K Population','Active Cases Per 100k Population','Total Test:Positive Ratio',
 'New Positive%','Case Fatality Rate%','New Confirmed Case Growth Rate','New Death Case Growth Rate','Average daily cases per 100,000 people in the past week',
 'New Test','NPI','Region','key-id','Country/District','7 days inc cases','7 days inc deaths']

raw_data=[]
for i in tqdm(range(start,end)):
    text_source=x[i]
    l=text_source.find_all('a',attrs={'class':'mt_a'})
    if l==[]:
        continue
    s=str(l[0])
    coun=s.split('/')
    try:
        region=cate[coun[1]]
        iso=dict_iso[coun[1]]
    except:
        region='missing'
    url='https://www.worldometers.info/coronavirus/country/'+coun[1]+'/'
    print(coun[1])
    a=''
    while a=='':
        try:
            a=requests.get(url,headers=headers)
        except:
            a=''
    soup = BeautifulSoup(a.content,'html5lib')
    r=soup.body.find_all('script',attrs={'type':'text/javascript'})
    p=re.compile(r'categories: \[(.*?)\]',re.S)
    rs=re.findall(p,r[0].text)
    d=rs[0]
    # d=re.sub(r'\"','',d)
    str_pat = re.compile(r'\"(.*?)\"')
    d = str_pat.findall(d)
    date=d
    p1=re.compile(r'name: \'Cases\'.*?\[(.*?)\]',re.S)
    for j in range(10):
        try:
            rs=re.findall(p1,r[j].text)
            d=rs[0]
            d=re.sub(r'\"','',d)
            case=d.split(',')
        except:
            # print('{} cases is not{}'.format(coun[1],j))
            continue
    p1=re.compile(r'name: \'Deaths\'.*?\[(.*?)\]',re.S)
    for j in range(10):
        try:
            rs=re.findall(p1,r[j].text)
            d=rs[0]
            d=re.sub(r'\"','',d)
            TD=d.split(',')
        except:
            continue
    j={'Date':date,'Total Cases':case,'Total Deaths':TD}
    # print(j)
    print("Date {} TC {} TD {}".format(len(date),len(case),len(TD)))
    if not len(set([len(date),len(case),len(TD)])) == 1:
        continue
    hist_data_of_coun_i=pd.DataFrame(j)
    hist_data_of_coun_i['Total Deaths'][0]=0
    for k in range(len(hist_data_of_coun_i['Total Deaths'])):
        if hist_data_of_coun_i['Total Deaths'][k]=='null':
            data['Total Deaths'][k]=0
    hist_data_of_coun_i['Total Cases']=hist_data_of_coun_i['Total Cases'].astype(int)
    hist_data_of_coun_i['Total Deaths']=hist_data_of_coun_i['Total Deaths'].astype(int)
    hist_data_of_coun_i['case inc']=hist_data_of_coun_i['Total Cases'].diff()
    hist_data_of_coun_i['death inc']=hist_data_of_coun_i['Total Deaths'].diff()
    #七日新增死亡与cases
    seven_cases=sum([hist_data_of_coun_i.loc[len(date)-i,'case inc'] for i in range(1,8)])
    seven_deaths=sum([hist_data_of_coun_i.loc[len(date)-i,'death inc'] for i in range(1,8)])

    inc1=hist_data_of_coun_i.loc[len(date)-1,'case inc']/(7*hist_data_of_coun_i.loc[len(date)-8,'case inc'])
    inc2=hist_data_of_coun_i.loc[len(date)-1,'death inc']/(7*hist_data_of_coun_i.loc[len(date)-8,'death inc'])

    inc_1=sum([hist_data_of_coun_i.loc[len(date)-i,'case inc'] for i in range(1,8)])/sum([hist_data_of_coun_i.loc[len(date)-i,'case inc'] for i in range(8,15)])
    inc_2=sum([hist_data_of_coun_i.loc[len(date)-i,'death inc'] for i in range(1,8)])/sum([hist_data_of_coun_i.loc[len(date)-i,'death inc'] for i in range(8,15)])
    adcp=sum([hist_data_of_coun_i.loc[len(date)-i,'case inc'] for i in range(1,8)])/7
    p=1
    while inc1 ==0 and hist_data_of_coun_i.loc[len(date)-1,'Total Cases']>=10000:
        p+=1
        inc1=hist_data_of_coun_i.loc[len(date)-p,'case inc']/(7*hist_data_of_coun_i.loc[len(date)-1-p,'case inc'])
    dd=hist_data_of_coun_i.shift(5)
    hist_data_of_coun_i['inc_p']=np.log(hist_data_of_coun_i['case inc']/dd['case inc'])/5
    hist_data_of_coun_i=hist_data_of_coun_i[~hist_data_of_coun_i.isin([np.nan, np.inf, -np.inf]).any(1)]
    da=hist_data_of_coun_i['inc_p'].values
    try:
        slope,intercept, r_value, p_value, std_err=st.linregress(list(range(30)), da[:30])
    except:
        slope=None

    # print(x[i])
    bo=x[i].text.split('\n')
    # print(bo)
    if bo[6]=='' and bo[7]=='':
        del bo[7]
    if bo[17]=='' and bo[18]=='':
        del bo[18]
    for o in range(start2,end2):
        s1=x[o]
        l1=s1.find_all('a',attrs={'class':'mt_a'})
        if l1==[]:
            continue
        s1=str(l1[0])
        coun1=s1.split('/')
        if coun1[1]==coun[1]:
            bo1=x[o].text.split('\n')
            break
    for h in range(len(bo)):
        bo[h]=bo[h].replace(',','')
        bo[h]=bo[h].replace('+','')
    for h in range(len(bo1)):
        bo1[h]=bo1[h].replace(',','')
        bo1[h]=bo1[h].replace('+','')
    #Cases Per 100K Population
    try:
        bo.append(100000*int(bo[3])/int(bo[15]))
    except:
        print(coun[1])
        continue
        # bo.append(np.nan)
        # print('lack one')
    #Tests Per 100K Population
    try:
        bo.append(100000*int(bo[13])/int(bo[15]))
    except:
        print(coun[1])
        continue
        # bo.append(np.nan)
        # print('lack one')
    #'Active Cases Per 100k Population'
    try:
        bo.append(int(bo[9])*100000/int(bo[15]))
    except:
        bo.append(np.nan)
        # print('lack one')

    #Total Test:Positive Ratio
    bo.append(int(bo[3])/int(bo[13]))
    #New Positive
    try:
        bo.append((int(bo[3])-int(bo1[3]))/(int(bo[13])-int(bo1[13])))
    except:
        bo.append(np.nan)
        # print('lack one')
    #Case Fatality Rate%
    try:
        if bo[5]=='':
            bo.append(0)
        else:
            bo.append(int(bo[5])/int(bo[3]))
    except:
        bo.append(np.nan)
    #New Confirmed Case Growth Rate
#     try:
#         q=2
#         while (math.isnan(inc1) or inc1==np.inf) and q<=9:
# #             print(inc1)
#             inc1=hist_data_of_coun_i.loc[len(date)-q,'case inc']/(7*hist_data_of_coun_i.loc[len(date)-q-7,'case inc'])
#             c=hist_data_of_coun_i.loc[len(date)-q,'case inc']
#             q+=1
# #         print(inc1)
#         if math.isnan(inc1):
#             bo.append(0)
#         elif inc1==np.inf:
#             bo.append(0.01)
#     #    elif c<=100:
#     #        bo.append(0.03)
#         else:
#             bo.append(inc1)
#     except:
#         bo.append(0)
        # print('lack one')
    #New Sum Confirmed Case Growth Rate
    if math.isnan(inc_1) or inc_1=='':
        bo.append(0)
    elif inc_1==np.inf:
        bo.append(0.01)
    else:
        bo.append(inc_1)
    print(bo[-1])
    #New Death Case Growth Rate
#     try:
#         q=2
#         while (math.isnan(inc2) or inc2==np.inf) and q<=9:
# #             print(inc2)
#             inc2=hist_data_of_coun_i.loc[len(date)-q,'death inc']/(7*hist_data_of_coun_i.loc[len(date)-q-7,'death inc'])
#             q+=1
# #         print(inc2)
#         if math.isnan(inc2):
#             bo.append(0)
#         elif inc2==np.inf:
#             bo.append(0.1)
#         else:
#             bo.append(inc2)
#     except:
#         bo.append(0)
        # print('lack one')
    #New Sum Death Case Growth Rate
    if math.isnan(inc_2) or inc_2=='':
        bo.append(0)
    elif inc_2==np.inf:
        bo.append(0.1)
    else:
        bo.append(inc_2)
    print(bo[-1])

    #Average daily cases per 100,000 people in the past week
    bo.append(adcp*100000/int(bo[15]))

    # New Test
    try:
        bo.append(int(bo[13])-int(bo1[13]))
    except:
        bo.append(np.nan)
        # print('lack one')

    bo.append(slope)
    if region=='missing':
        continue
    else:
        bo.append(region)
    bo.append(coun1[1])
    bo.append(iso)
    bo.append(seven_cases)
    bo.append(seven_deaths)
    print(len(bo))
    print(bo)
    if len(bo)!=39:
        os.exit(-1)
    raw_data.append(bo)
raw_data=DataFrame(raw_data,columns=col_name)
brief_raw_data=raw_data[['Country,Other','key-id','Region','Country/District','Population',
'TotalCases','ActiveCases','TotalDeaths','NewDeaths','TotalRecovered','NewRecovered','Serious,Critical','NewCases','New Test','Cases Per 100K Population','Tests Per 100K Population',
'Active Cases Per 100k Population','Total Test:Positive Ratio','New Positive%',
          'Case Fatality Rate%','New Confirmed Case Growth Rate','New Death Case Growth Rate','Average daily cases per 100,000 people in the past week','NPI','7 days inc cases','7 days inc deaths']]
brief_raw_data['week death rate']=brief_raw_data['7 days inc deaths']/brief_raw_data['7 days inc cases']
tf=copy.deepcopy(brief_raw_data)
uni_region=list(set(list(tf['Region'].values)))

uni_region.remove('western europe')
data_region=tf[tf['Region']=='western europe']
data_region=data_region.replace(np.nan,'shit')
data_region=data_region.replace(np.inf,'shit')
data_region=data_region.replace('N/A','shit')
data_region=data_region.replace('',0)
data_region=data_region.replace(' ',0)
data_region.loc[data_region['NPI']=='shit','NPI']=0
data_region.loc[data_region['Case Fatality Rate%']=='shit','Case Fatality Rate%']=0
dd=data_region[['TotalCases','ActiveCases','TotalRecovered','Case Fatality Rate%']]
ac=dd[(dd['TotalCases']!='shit')&(dd['ActiveCases']!='shit')&(dd['TotalRecovered']!='shit')&(dd['Case Fatality Rate%']!='shit')]
active_rate_region=sum(ac['ActiveCases'].astype(int))/sum(ac['TotalCases'].astype(int))
data_region.loc[data_region['Active Cases Per 100k Population']=='shit','Active Cases Per 100k Population']=active_rate_region*100000*data_region.loc[data_region['Active Cases Per 100k Population']=='shit','TotalCases'].astype(int)/data_region.loc[data_region['Active Cases Per 100k Population']=='shit','Population'].astype(int)

dd=data_region[['NewCases','New Test']]
ac=dd[dd['New Test']!=0]
new_posi=sum(ac['NewCases'].astype(int))/sum(ac['New Test'])
data_region.loc[data_region['New Test']==0,'New Positive%']=new_posi
final=copy.deepcopy(data_region)
for distri in uni_region:
    data_region=tf[tf['Region']==distri]
    data_region=data_region.replace(np.nan,'shit')
    data_region=data_region.replace(np.inf,'shit')
    data_region=data_region.replace('N/A','shit')
    data_region=data_region.replace('',0)
    data_region=data_region.replace(' ',0)
    data_region.loc[data_region['NPI']=='shit','NPI']=0
    data_region.loc[data_region['Case Fatality Rate%']=='shit','Case Fatality Rate%']=0
    dd=data_region[['TotalCases','ActiveCases','TotalRecovered','Case Fatality Rate%']]
    ac=dd[(dd['TotalCases']!='shit')&(dd['ActiveCases']!='shit')&(dd['TotalRecovered']!='shit')&(dd['Case Fatality Rate%']!='shit')]
    active_rate_region=sum(ac['ActiveCases'].astype(int))/sum(ac['TotalCases'].astype(int))
    data_region.loc[data_region['Active Cases Per 100k Population']=='shit','Active Cases Per 100k Population']=active_rate_region*100000*data_region.loc[data_region['Active Cases Per 100k Population']=='shit','TotalCases'].astype(int)/data_region.loc[data_region['Active Cases Per 100k Population']=='shit','Population'].astype(int)

    dd=data_region[['NewCases','New Test']]
    ac=dd[dd['New Test']!=0]
    try:
        new_posi=sum(ac['NewCases'].astype(int))/sum(ac['New Test'])
    except:
        new_posi=0
    data_region.loc[data_region['New Test']==0,'New Positive%']=new_posi
    data_region.loc[data_region['New Test']=='shit','New Positive%']=new_posi
    final=pd.concat([final,data_region])
final=final.reset_index(drop=True)
tf2=final[['Country,Other','key-id','Country/District','Region','TotalCases','Cases Per 100K Population','Tests Per 100K Population',
'Active Cases Per 100k Population','Total Test:Positive Ratio','New Positive%',
          'Case Fatality Rate%','New Confirmed Case Growth Rate','New Death Case Growth Rate','Average daily cases per 100,000 people in the past week','NPI']]

#越高越好，即需要降序
# for x in ['Cases Per 100K Population','Active Cases Per 100k Population','Total Test:Positive Ratio','New Positive%',]
x='Tests Per 100K Population'
df=copy.deepcopy(tf2[['Country,Other',x]])
df2=df.sort_values(x,ascending=False,inplace=False)
df2 = df2.reset_index(drop=True)
df2['cum']=df.index+1
df2['cum_prob']=100*df2['cum']/max(df2['cum'])
df3=pd.merge(df,df2,on=['Country,Other'])
# tf2['IND_'+x]=df3['cum_prob']
tf2['IND_'+x]=0
for h in list(tf2['Country,Other'].values):
    tf2.loc[tf2['Country,Other']==h,'IND_'+x]=df3.loc[df3['Country,Other']==h,'cum_prob'].values[0]

for x in ['Cases Per 100K Population','Active Cases Per 100k Population','Total Test:Positive Ratio','New Positive%','Case Fatality Rate%','New Confirmed Case Growth Rate','New Death Case Growth Rate','Average daily cases per 100,000 people in the past week','NPI']:
    i=1
    df=copy.deepcopy(tf2[['Country,Other',x]])
    print(x)
    df2=df.sort_values(x,inplace=False)
    df2 = df2.reset_index(drop=True)
    df2['cum']=df.index+1
    df2['cum_prob']=100*df2['cum']/max(df2['cum'])
    df3=pd.merge(df,df2,on=['Country,Other'])
    tf2['IND_'+x]=0
    for h in list(tf2['Country,Other'].values):
        tf2.loc[tf2['Country,Other']==h,'IND_'+x]=df3.loc[df3['Country,Other']==h,'cum_prob'].values[0]
    i+=1

tf2['Comprehensive Index']=0.15*tf2['IND_Cases Per 100K Population']+0.08*tf2['IND_Tests Per 100K Population']+0.2*tf2['IND_Active Cases Per 100k Population']+0.1*tf2['IND_Total Test:Positive Ratio']+0.13*tf2['IND_New Positive%']+0.05*tf2['IND_Case Fatality Rate%']+ 0.22*tf2['IND_New Confirmed Case Growth Rate']+0.07*tf2['IND_New Death Case Growth Rate']
today=datetime.datetime.now()
tf4=tf2[['Country/District','TotalCases','IND_Cases Per 100K Population','IND_Tests Per 100K Population','IND_Total Test:Positive Ratio',
'IND_New Positive%','IND_Case Fatality Rate%','IND_New Confirmed Case Growth Rate','IND_New Death Case Growth Rate','IND_Active Cases Per 100k Population',
'IND_NPI','IND_Average daily cases per 100,000 people in the past week','Comprehensive Index']]
tf_c=copy.deepcopy(tf4)
tf_c_rename=tf_c.rename({'TotalCases':'TOTAL CASE','IND_Cases Per 100K Population':'IND1_Cases Per 100K Population','IND_Tests Per 100K Population':'IND2_Tests Per 100K Population',
'IND_Active Cases Per 100k Population':'IND8_Active Cases Per 100k Population','IND_Total Test:Positive Ratio':'IND3_Total Test:Positive Ratio',
'IND_New Positive%':'IND4_New Positive%','IND_Case Fatality Rate%':'IND5_Case Fatality Rate%','IND_New Confirmed Case Growth Rate':'IND6_New Confirmed Case Growth Rate',
'IND_New Death Case Growth Rate':'IND7_New Death Case Growth Rate','IND_NPI':'NPI'},axis='columns')
tf_c_rename.to_excel('World_index_{}.xlsx'.format(today),sheet_name=ts[1],index=False)
tf2.to_excel('World_raw_index_{}.xlsx'.format(today),sheet_name=ts[1],index=False)

url='https://raw.githubusercontent.com/owid/covid-19-data/master/public/data/vaccinations/vaccinations.csv'
a=requests.get(url,headers=headers)
with open("vacc.csv",'wb') as f:
    f.write(a.content)
vacc = pd.read_csv('vacc.csv',keep_default_na=False)
ct = list(dict(vacc['location'].value_counts()).keys())
for x in ['total_vaccinations','people_vaccinated','people_fully_vaccinated','total_boosters']:
    vacc[x]=vacc[x].replace('',0)
    vacc[x]=vacc[x].astype(float)
    vacc[x]=vacc[x].astype(int)
img = dict()
for i in ct:
    dt = vacc[vacc['location']==i]
    d=[]
    for x in ['total_vaccinations','people_vaccinated','people_fully_vaccinated','total_boosters']:
        d.append(max(dt[x]))
    img[i]=d
brief_raw_data['total_vaccinations']=0
brief_raw_data['people_vaccinated']=0
brief_raw_data['people_fully_vaccinated']=0
brief_raw_data['total_boosters']=0
for i in img.keys():
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'total_vaccinations'] = int(img[i][0])
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'people_vaccinated'] = int(img[i][1])
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'people_fully_vaccinated'] = int(img[i][2])
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'total_boosters'] = int(img[i][3])

brief_raw_data['Population']=brief_raw_data['Population'].astype(int)
brief_raw_data['vacc_per_100']=brief_raw_data['total_vaccinations']*100/brief_raw_data['Population']
brief_raw_data['cases_per_100']=brief_raw_data['Cases Per 100K Population']/1000
brief_raw_data['total_immune']=brief_raw_data['cases_per_100']+brief_raw_data['vacc_per_100']*0.9

def infer_vaccinated(one_row):
    if one_row['total_vaccinations']!=0 and one_row['people_vaccinated']==0:
        one_row['people_vaccinated']=int(0.7018*one_row['total_vaccinations'])
    if one_row['total_vaccinations']!=0 and one_row['people_fully_vaccinated']==0:
        one_row['people_fully_vaccinated']=int(0.2936*one_row['total_vaccinations'])
    return one_row
brief_raw_data = brief_raw_data.apply(infer_vaccinated,axis=1)

import datetime
from datetime import timedelta
import sys
vacc = pd.read_csv('vacc.csv',keep_default_na=False)
ct = list(dict(vacc['location'].value_counts()).keys())
vacc_daily=copy.deepcopy(vacc)
vacc_daily['daily_vaccinations']=vacc_daily['daily_vaccinations'].replace('',-1)
vacc_daily['daily_vaccinations']=vacc_daily['daily_vaccinations'].astype(int)
vacc_daily = vacc_daily.drop(vacc_daily[vacc_daily['daily_vaccinations']==-1].index)
vacc_daily['date']=pd.to_datetime(vacc_daily['date']).dt.date
daily=dict()
# vacc_daily.loc[vacc_daily['date']==,'daily_vaccinations']
for i in ct:
    rdd=[]
    dt = vacc_daily[vacc_daily['location']==i]
    today=datetime.date.today()
    try:
        a=max(dt['date'])
    #     print(a)
    #     print(today)
    #     print(today-timedelta(14))
#         if today-timedelta(14)<a:
        x=dt.loc[(dt['date']<=today)&(dt['date']>=today-timedelta(14)),'daily_vaccinations']
        note = sum(x)/len(x)
#         else:
#             note = -2
    except:
        note =-2
    daily[i]=int(note)
for x in ['total_vaccinations','people_vaccinated','people_fully_vaccinated','total_boosters']:
    vacc[x]=vacc[x].replace('',-1)
    vacc[x]=vacc[x].astype(float)
    vacc[x]=vacc[x].astype(int)
vacc = vacc.drop(vacc[vacc['total_vaccinations']==-1].index)
vacc['date']=pd.to_datetime(vacc['date'])
rd=dict()
for i in ct:
    rdd=[]
    dt = vacc[vacc['location']==i]
    a=max(dt['date'])
    b=min(dt['date'])
    c=a-b
    Vacc_Days=int(c.days)
    if Vacc_Days==0:
        Vacc_Days=1
    Total_Vacc=int(max(dt['total_vaccinations']))
    Avg_Vac_Daily=Total_Vacc/Vacc_Days
    d=a-timedelta(14)
    if len(list(dt['date']))==1:
        Avg_Vac_Last14D=Total_Vacc
    elif daily[i]==-2:
        Avg_Vac_Last14D=Total_Vacc/Vacc_Days
    else:
        Avg_Vac_Last14D=daily[i]
    rdd=[Vacc_Days,Total_Vacc,Avg_Vac_Daily,Avg_Vac_Last14D]
    rd[i]=rdd

# lst14days=pd.DataFrame(rd,columns=['Country','Vac_Days','Total_Vac','Avg_Vac_Daily','Avg_Vac_Last14D'])
brief_raw_data['累计感染率（%）']=brief_raw_data['Cases Per 100K Population']/100000
brief_raw_data['接种率=（接种疫苗数/总人口）']=brief_raw_data['vacc_per_100']/100
brief_raw_data['VacSpeed_Last14D']=0.05
brief_raw_data['Population']=brief_raw_data['Population'].astype(int)

brief_raw_data['fully vaccinated %']=100*brief_raw_data['people_fully_vaccinated']/brief_raw_data['Population']
brief_raw_data['at least 1 dose vaccinated %']=100*brief_raw_data['people_vaccinated']/brief_raw_data['Population']
brief_raw_data['Fully Vac Ratio %']=100*brief_raw_data['fully vaccinated %']/brief_raw_data['at least 1 dose vaccinated %']
# brief_raw_data['加强针接种率']=brief_raw_data['total_boosters']/brief_raw_data['Population']

for i in rd.keys():
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'VacSpeed_Last14D'] = 100*int(rd[i][3])/brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'Population']
vaccine_efficacy = pd.read_excel('Vaccine-Country.xlsx',sheet_name='Efficacy')
vaccine_efficacy = vaccine_efficacy.set_index('Vaccine').T.to_dict('list')
vaccine_map=pd.read_excel('Vaccine-Country.xlsx',sheet_name='Vac-Mapping')
vaccine_map = dict(zip(vaccine_map['Location'],vaccine_map['首选疫苗']))
brief_raw_data['Vaccine1']='NA'
brief_raw_data['Vaccine_Efficacy']=0.8
for i in vaccine_map.keys():
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'Vaccine1'] = vaccine_map[i]
    brief_raw_data.loc[((brief_raw_data['Country,Other']==i)|(brief_raw_data['Country/District']==i)),'Vaccine_Efficacy'] = vaccine_efficacy[vaccine_map[i]][1]

brief_raw_data['total_immune_adj']=brief_raw_data['累计感染率（%）']*100+brief_raw_data['vacc_per_100']*brief_raw_data['Vaccine_Efficacy']/2
brief_raw_data['DaystoReachHerdImmunity']=(0.7-brief_raw_data['累计感染率（%）']-brief_raw_data['接种率=（接种疫苗数/总人口）']*brief_raw_data['Vaccine_Efficacy']/2)/(brief_raw_data['VacSpeed_Last14D']*brief_raw_data['Vaccine_Efficacy']*0.01/2)

import datetime
def change_time(x):
    in_date = ts[1]
    # print(x)
    if x<0:
        out_date = 'already'
        return out_date
    elif math.isinf(x):
        out_date = 'never'
        return out_date
    dt = datetime.datetime.strptime(in_date, "%Y%m%d")
    try:
        out_date = (dt + datetime.timedelta(days=int(x))).strftime("%Y-%m-%d")
    except:
        out_date = 'never'
    return out_date

# brief_raw_data.to_excel('World_rawdata_test.xlsx')
brief_raw_data['HerdImmunityDate']=brief_raw_data['DaystoReachHerdImmunity'].map(change_time)
brief_raw_data['Total Boosters']=brief_raw_data['total_boosters']
brief_raw_data['加强针接种率']=brief_raw_data['total_boosters']/brief_raw_data['Population']
brief_raw_data.to_excel('World_rawdata_{}.xlsx'.format(today),sheet_name=ts[1],index=False)

uni_region=list(set(list(brief_raw_data['Region'].values)))
writer = pd.ExcelWriter('Region_World_Rawdata_{}.xlsx'.format(today))
for distri in uni_region:
    data_region=brief_raw_data[brief_raw_data['Region']==distri]
    data_region.to_excel(writer,sheet_name=distri,index=False)
writer.save()


#!/usr/bin/env python

"""
Plot distances between profiles.

"""
from __future__ import print_function

from optparse import OptionParser
import numpy as np
from scipy.spatial.distance import cdist
import matplotlib.pyplot as plt
import pylab
import cpa
import sys
from .profiles import Profiles

def plot_distances(profiles, output_group_name=None):
    if output_group_name:
        input_group_r, input_colnames = cpa.db.group_map(profiles.group_name,
                                                         reverse=True)
        input_group_r = dict((tuple(map(str, k)), v)
                             for k, v in input_group_r.items())
        output_group, output_colnames = cpa.db.group_map(output_group_name)
        d = {}
        labels = []
        for i, k in enumerate(profiles.keys()):
            groups = [output_group[image] for image in input_group_r[k]]
            if groups.count(groups[0]) != len(groups):
                print('Error: Input group %r contains images in %d output groups' % (key, len(set(groups))), file=sys.stderr)
                sys.exit(1)
            d.setdefault(groups[0], []).append(i)
            labels.append(groups[0])
        ordering = [i for k in sorted(d.keys()) for i in d[k]]
        labels = list(np.array(labels)[ordering])
    else:
        ordering = np.arange(len(profiles.keys))
        labels = list(np.array(profiles.keys())[ordering])
    labels = [' '.join(map(str, k)) for k in labels]
    data = profiles.data[ordering]

    for i in range(len(labels))[:0:-1]:
        if labels[i] == labels[i - 1]:
            labels[i] = ''

    fig, ax = plt.subplots()
    dist = cdist(data, data, 'cosine')
    axes_image = ax.imshow(dist, cmap=pylab.cm.RdBu, interpolation='nearest')
    fig.colorbar(axes_image, use_gridspec=True)
    ax.set_yticks([i for i, l in enumerate(labels) if l != ''])
    ax.set_yticklabels([l for l in labels if l != ''])
    for tick in ax.yaxis.iter_ticks():
        tick[0].label1.set_fontsize(8)
    ax.set_xticks([i for i, l in enumerate(labels) if l != ''])
    ax.set_xticklabels(['' for l in labels if l != ''])
    plt.axis('image')
    plt.tight_layout()

def parse_arguments():
    parser = OptionParser("usage: %prog PROPERTIES-FILE INPUT-FILENAME GROUP")
    parser.add_option('-o', dest='output_filename', help='file to store the profiles in')
    options, args = parser.parse_args()
    if len(args) != 3:
        parser.error('Incorrect number of arguments')
    return options, args

if __name__ == '__main__':
    options, (properties_file, input_filename, group_name) = parse_arguments()
    cpa.properties.LoadFile(properties_file)
    profiles = Profiles.load(input_filename)
    plot_distances(profiles, group_name)
    if options.output_filename:
        pylab.savefig(options.output_filename)
    else:
        pylab.show()


<reponame>markflorisson/minivect
"""
Determine data dependences between two arbitrary numpy arrays.
See thesis/thesis.pdf section 5.2.
"""

import fractions

import numpy as np

a = np.empty((10, 10))

def get_start_end(a):
    start = end = a.ctypes.data
    for i in range(a.ndim):
        if a.shape[i] == 0:
            # Empty array, no overlap possible
            return 0, 0
        else:
            offset = a.strides[i] * (a.shape[i] - 1)
            if a.strides[i] > 0:
                end += offset
            else:
                start += offset

    return start, end

def overlap(a, b):
    """
    >>> overlap(a, a)
    True
    >>> overlap(a[:5], a[5:])
    False
    >>> overlap(a[:6], a[5:])
    True
    >>> overlap(a[2:, 5:], a[5:, :6])
    True
    >>> overlap(a[2:3, 5:6], a[2:, 6:])
    False
    >>> overlap(a[2:3, 5:6], a[2:, 5:])
    True
    """
    start1, end1 = get_start_end(a)
    start2, end2 = get_start_end(b)

    if start1 < start2:
        return end1 > start2
    else:
        return end2 > start1

def verify_containment(base, a):
    base_start, base_end = get_start_end(base)
    a_start, a_end = get_start_end(a)
    assert a_start >= base_start and a_end <= base_end

def verify_base(base, a, b):
    # verify ndim
    assert base.ndim == a.ndim

    # verify pointers
    verify_containment(base, a)
    verify_containment(base, b)

    # verify strides
    negative = base.strides[0] < 0
    for i in range(base.ndim):
        # base strides may not be 0, and must be either all positive, or all negative
        assert base.strides[i] != 0
        assert (base.strides[i] < 0) == negative

        assert a.strides[i] % base.strides[i] == 0, (a.strides[i], base.strides[i])
        assert b.strides[i] % base.strides[i] == 0, (b.strides[i], base.strides[i])

def array_order(a, is_base=False):
    """
    Return the dimension indices corresponding to ascending stride order

    >>> a = np.empty((10, 10, 10), order='C')
    >>> array_order(a)
    [2, 1, 0]
    >>> array_order(a[::-1, ::-1])
    [2, 1, 0]
    >>> array_order(a.T)
    [0, 1, 2]
    >>> array_order(a.swapaxes(0, 1))
    [2, 0, 1]
    """
    stride_index_pairs = zip((abs(stride) for stride in a.strides), range(a.ndim))
    stride_index_pairs.sort()
    order = [index for stride, index in stride_index_pairs]
    if is_base and a.strides[0] < 0:
        return order[::-1]
    return order

def dimensional_independence(a):
    """
    >>> dimensional_independence(a)
    >>> dimensional_independence(a[:1, ::2])
    >>> dimensional_independence(a[::-2, ::-2])
    >>> d = a.copy()
    >>> d.strides = (10, 2) # 2 * 10 = 20, so columns overlap with rows
    >>> dimensional_independence(d)
    Traceback (most recent call last):
       ...
    AssertionError: (10, 18)
    """
    order = array_order(a)
    for previous_dim, dim in zip(order, order[1:]):
        extent = a.strides[previous_dim] * (a.shape[previous_dim] - 1)
        assert abs(a.strides[dim]) >= abs(extent), (a.strides[dim], extent)

def verify_dimensional_constraints(base, a, offsets, steps):
    for i, (offset, step) in enumerate(zip(offsets, steps)):
        assert 0 <= offset + (a.shape[i] - 1) * step < base.shape[i]

def offsets(base, a):
    """
    >>> offsets(a, a[2:, 4:])
    (2, 4)
    >>> offsets(a, a[::-1, ::-1])
    (9, 9)
    >>> offsets(a, a[::-1])
    (9, 0)
    >>> offsets(a, a[:, ::-1])
    (0, 9)
    """
    # assert array_order(base) == array_order(a)
    offsets = [None] * base.ndim
    distance = a.ctypes.data - base.ctypes.data
    for dim in reversed(array_order(base, is_base=True)):
        offsets[dim] = distance // base.strides[dim]
        assert 0 <= offsets[dim] < base.shape[dim], (offsets[dim], base.shape[dim])
        distance = distance % base.strides[dim]

    assert distance == 0

    return tuple(offsets)

def steps(base, a):
    steps = []
    for i in range(base.ndim):
        steps.append(a.strides[i] // base.strides[i])
    return steps

def verify_solution(base, a, offsets_a, steps_a):
    slices = []
    for offset, step, extent in zip(offsets_a, steps_a, a.shape):
        slices.append(slice(offset, offset + extent * step, step))

    np.all(base[tuple(slices)] == a)

def verify(base, a, b):
    """
    >>> verify(a, a[::2, ::2], a[1::2, 1::2])
    ((0, 0), [2, 2], (1, 1), [2, 2])
    >>> verify(a, a, a[::-1, ::-1])
    ((0, 0), [1, 1], (9, 9), [-1, -1])
    >>> verify(a, a, a[:, ::-1])
    ((0, 0), [1, 1], (0, 9), [1, -1])
    >>> verify(a, a[:, ::-1], a[::-1, :])
    ((0, 9), [1, -1], (9, 0), [-1, 1])
    >>> verify(a, a[:, 2::2], a[:, 8:0:-2])
    ((0, 2), [1, 2], (0, 8), [1, -2])
    """
    assert a.shape == b.shape

    verify_base(base, a, b)

    dimensional_independence(base)
    dimensional_independence(a)
    dimensional_independence(b)

    offsets_a = offsets(base, a)
    steps_a = steps(base, a)
    verify_dimensional_constraints(base, a, offsets_a, steps_a)

    offsets_b = offsets(base, b)
    steps_b = steps(base, b)
    verify_dimensional_constraints(base, b, offsets_b, steps_b)

    verify_solution(base, a, offsets_a, steps_a)
    verify_solution(base, b, offsets_b, steps_b)

    return (offsets_a, steps_a, offsets_b, steps_b)

def siv(base, a, b, offsets_a, steps_a, offsets_b, steps_b):
    direction_vector = []
    for offset_a, offset_b, step_a, step_b, extent in zip(offsets_a, offsets_b,
                                                          steps_a, steps_b, a.shape):
        distance = offset_b - offset_a
        if step_a == step_b:
            # SIV test
            if distance % step_a == 0 and abs(distance) < extent:
                # dependence
                if distance < 0:
                    direction_vector.append('>')
                elif distance == 0:
                    direction_vector.append('=')
                else:
                    direction_vector.append('<')
            else:
                # independence
                return None
        elif distance % fractions.gcd(step_a, step_b) != 0:
            # independence
            return None
        elif abs(step_a) == abs(step_b):
            # Weak crossing SIV test
            i = distance / (2*step_a)
            if i % 0.5 == 0 and 0 <= i < extent:
                # dependence
                direction_vector.append('*')
            else:
                # independence
                return None
        else:
            # possible dependence, further checking is needed (try banerjee). Remain conservative
            direction_vector.append('*')

    return tuple(direction_vector)

def direction_vector(a, b):
    """
    Returns False in case of independence, or a direction vector if no
    independence is proven

    >>> [direction_vector(*slice_pair) for slice_pair in slice_pairs]
    [False, False, False, False, False, False, ('*', '=')]

    # Dependent, test direction vectors
    >>> direction_vector(a[::2], a[8::-2])
    ('*', '=')
    >>> direction_vector(a[::2, 1:5], a[8::-2, 0:4])
    ('*', '>')
    >>> direction_vector(a[::2, 1:5], a[8::-2, 2:6])
    ('*', '<')
    """
    if not overlap(a, b):
        return False

    assert a.shape == b.shape

    base = a
    while base.base is not None:
        base = base.base

    base_b = b
    while base_b.base is not None:
        base_b = base_b.base

    if not (base.ctypes.data == base_b.ctypes.data and
            base.shape == base_b.shape and
            base.strides == base_b.strides):
        return True # assume dependence

    offsets_steps = verify(base, a, b)

    result = siv(base, a, b, *offsets_steps)
    if result is None:
        return False # no dependence

    return result # direction vector

def general_gcd_test(a, b):
    """
    Returns False in case of independence, otherwise there may or may not
    be a dependence

    >>> [general_gcd_test(*slice_pair) for slice_pair in slice_pairs]
    [False, True, False, False, True, True, True]
    """
    # perform the GCD test without prior transformation
    if overlap(a, b):
        gcd = fractions.gcd(reduce(fractions.gcd, a.strides),
                            reduce(fractions.gcd, a.strides))
        return (a.ctypes.data - b.ctypes.data) % gcd == 0

    return False

slice_pairs = [
    # independent
    (a[::2, ::2], a[1::2, 1::2]),
    (a[::2], a[::-2]),
    (a[:, ::2], a[:, ::-2]),
    (a[:5], a[5:]),
    (a[:, :5], a[:, 5:]),
    (a[1:3, 4:8], a[2:4, 0:4]),
    (a[0:4:3], a[2:5:2]), # rows (0, 3) and (2, 4)
]

if __name__ == '__main__':
    import doctest
    doctest.testmod()


<filename>crossval_hyper.py
#Python2,3 compatible headers
from __future__ import unicode_literals,division
from builtins import int
from builtins import range

#System packages
import torch
from torch.autograd import Variable,grad
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import numpy
import scipy
import scipy.misc
import math
import time
import random
import argparse
import sys
import os
import re
import copy
import importlib
from collections import namedtuple
from collections import OrderedDict
from itertools import chain

import PIL.Image
import torchvision.datasets.folder
import torchvision.transforms.functional as Ft
import torchvision.transforms as Ts
import PIL.Image as Image

import torch.utils.data.dataloader as myDataLoader
import skimage.io

import util.db as db
import util.smartparse as smartparse
import util.file
import util.session_manager as session_manager
import dataloader

import sklearn.metrics
from hyperopt import hp, tpe, fmin

# Training settings
def default_params():
    params=smartparse.obj();
    #Data
    params.nsplits=5;
    params.pct=0.8
    #Model
    params.arch='arch.mlpv2';
    params.nh=256;
    params.nlayers=3;

    #Optimization
    params.batch=256;
    params.lr=1e-3;
    params.epochs=300;
    params.decay=1e-4;

    #MISC
    params.session_dir=None;
    params.budget=10000;
    return params

def create_session(params):
    session=session_manager.Session(session_dir=params.session_dir); #Create session
    torch.save({'params':params},session.file('params.pt'));
    pmvs=vars(params);
    pmvs=dict([(k,pmvs[k]) for k in pmvs if not(k=='stuff')]);
    print(pmvs);
    util.file.write_json(session.file('params.json'),pmvs); #Write a human-readable parameter json
    session.file('model','dummy');
    return session;

params = smartparse.parse()
params = smartparse.merge(params, default_params())
params.argv=sys.argv;

data=dataloader.new('data_r4v2.pt');
data.cuda();
params.stuff=data.preprocess();

session=create_session(params);
params.session=session;

#Hyperparam search config
hp_config=[];

#   Architectures
#archs=['arch.mlpv2','arch.mlpv3','arch.mlpv4','arch.mlpv5','arch.mlpv6'];
archs=['arch.mlp_set_color_v2xy'];

hp_config.append(hp.choice('arch',archs));
hp_config.append(hp.qloguniform('nh',low=math.log(16),high=math.log(384),q=4));
hp_config.append(hp.qloguniform('nh2',low=math.log(16),high=math.log(384),q=4));
hp_config.append(hp.quniform('nlayers',low=1,high=10,q=1));
hp_config.append(hp.quniform('nlayers2',low=1,high=10,q=1));
hp_config.append(hp.quniform('nlayers3',low=1,high=10,q=1));
#   OPT
hp_config.append(hp.qloguniform('epochs',low=math.log(5),high=math.log(200),q=1));
hp_config.append(hp.loguniform('lr',low=math.log(1e-4),high=math.log(5e-2)));
hp_config.append(hp.loguniform('decay',low=math.log(1e-6),high=math.log(1e-3)));
hp_config.append(hp.qloguniform('batch',low=math.log(16),high=math.log(64),q=1));

#Function to compute performance
def configure_pipeline(params,arch,nh,nh2,nlayers,nlayers2,nlayers3,epochs,lr,decay,batch):
    params_=smartparse.obj();
    params_.arch=arch;
    params_.nh=int(nh);
    params_.nh2=int(nh2);
    params_.nlayers=int(nlayers);
    params_.nlayers2=int(nlayers2);
    params_.nlayers3=int(nlayers3);
    params_.epochs=epochs;
    params_.lr=lr;
    params_.batch=batch;
    params_=smartparse.merge(params_,params);
    return params_;

crossval_splits=[];
for i in range(params.nsplits):
    data_train,data_test=data.generate_random_crossval_split(pct=params.pct);
    data_val,data_test=data_test.generate_random_crossval_split(pct=0.5);
    crossval_splits.append((data_train,data_val,data_test));

best_loss_so_far=-1e10;
def run_crossval(p):
    global best_loss_so_far
    max_batch=16;
    arch,nh,nh2,nlayers,nlayers2,nlayers3,epochs,lr,decay,batch=p;
    params_=configure_pipeline(params,arch,nh,nh2,nlayers,nlayers2,nlayers3,epochs,lr,decay,batch);
    arch_=importlib.import_module(params_.arch);
    #Random splits N times
    auc=[];
    ce=[];
    cepre=[];
    results_by_key={};
    t0=time.time();
    ensemble=[];
    for split_id,split in enumerate(crossval_splits):
        data_train,data_val,data_test=split;
        net=arch_.new(params_).cuda();
        opt=optim.Adam(net.parameters(),lr=params_.lr); #params_.lr

        #Train loop
        best_loss=-1e10;
        best_net=copy.deepcopy(net);

        for iter in range(params_.epochs):
            net.train();
            loss_total=[];
            for data_batch in data_train.batches(params_.batch,shuffle=True):
                opt.zero_grad();
                net.zero_grad();
                data_batch.cuda();
                C=data_batch['label'];
                data_batch.delete_column('label');
                scores_i=net(data_batch);

                #loss=F.binary_cross_entropy_with_logits(scores_i,C.float());
                spos=scores_i.gather(1,C.view(-1,1)).mean();
                sneg=torch.exp(scores_i).mean();
                loss=-(spos-sneg+1);
                #print(float(loss))
                l2=0;
                for p in net.parameters():
                    l2=l2+(p**2).sum();

                loss=loss+l2*params_.decay;
                loss.backward();
                loss_total.append(float(loss));
                opt.step();

            loss_total=sum(loss_total)/len(loss_total);

            #Eval every epoch
            #net.eval();
            #scores=[];
            #gt=[]
            #for data_batch in data_val.batches(max_batch):
            #    data_batch.cuda();
            #
            #    C=data_batch['label'];
            #    data_batch.delete_column('label');
            #    scores_i=net(data_batch);
            #    scores.append(scores_i.data.cpu());
            #    gt.append(C.data.cpu());

            #scores=torch.cat(scores,dim=0);
            #gt=torch.cat(gt,dim=0);

            #auc_i=sklearn.metrics.roc_auc_score(gt.numpy(),scores.numpy());
            #loss_i=float(F.binary_cross_entropy_with_logits(scores,gt.float()));
            #if best_loss<auc_i:
            #    best_loss=auc_i;
            #    best_net=copy.deepcopy(net);

            #print('train %.4f, loss %.4f, auc %.4f'%(float(loss_total),float(loss_i),float(auc_i)))

        #Temperature-scaling calibration on val
        #net=best_net;
        net.eval();
        scores=[];
        gt=[];
        for data_batch in data_val.batches(max_batch):
            data_batch.cuda();

            C=data_batch['label'];
            data_batch.delete_column('label');
            scores_i=net.logp(data_batch);
            scores.append(scores_i.data);
            gt.append(C);

        scores=torch.cat(scores,dim=0);
        gt=torch.cat(gt,dim=0);

        T=torch.Tensor(1).fill_(0).cuda();
        T.requires_grad_();
        opt2=optim.Adamax([T],lr=3e-2);
        for iter in range(500):
            opt2.zero_grad();
            loss=F.binary_cross_entropy_with_logits(scores*torch.exp(-T),gt.float());
            loss.backward();
            opt2.step();

        #Eval
        net.eval();
        scores=[];
        scores_pre=[];
        gt=[]
        for data_batch in data_test.batches(max_batch):
            data_batch.cuda();

            C=data_batch['label'];
            data_batch.delete_column('label');
            scores_i=net.logp(data_batch);

            scores.append((scores_i*torch.exp(-T)).data.cpu());
            scores_pre.append(scores_i.data.cpu());

            gt.append(C.data.cpu());

        scores=torch.cat(scores,dim=0);
        scores_pre=torch.cat(scores_pre,dim=0);
        gt=torch.cat(gt,dim=0);

        def compute_metrics(scores,gt,keys=None):
            if not keys is None:
                #slicing
                categories=set([k for k in keys if not k is None]);
                results={};
                for c in categories:
                    ind=[i for i,k in enumerate(keys) if k==c or k is None];
                    scores_c=[scores[i] for i in ind];
                    gt_c=[gt[i] for i in ind];
                    auc_c,ce_c=compute_metrics(scores_c,gt_c);
                    results[c]={'auc':auc_c,'ce':ce_c};

                return results;
            else:
                #Overall
                auc=float(sklearn.metrics.roc_auc_score(torch.LongTensor(gt).numpy(),torch.Tensor(scores).numpy()));
                ce=float(F.binary_cross_entropy_with_logits(torch.Tensor(scores),torch.Tensor(gt)));
                return auc,ce;

        auc_i,ce_i=compute_metrics(scores.tolist(),gt.tolist());
        _,ce_pre_i=compute_metrics(scores_pre.tolist(),gt.tolist());
        results_i=compute_metrics(scores.tolist(),gt.tolist(),data_test.data['table_ann']['nclasses'])
        for k in results_i:
            if k in results_by_key:
                results_by_key[k]['auc'].append(results_i[k]['auc'])
                results_by_key[k]['ce'].append(results_i[k]['ce']);
            else:
                results_by_key[k]={'auc':[],'ce':[]};

        #results_i=compute_metrics(scores.tolist(),gt.tolist(),data_test.data['table_ann']['arch'])
        #for k in results_i:
        #    if k in results_by_key:
        #        results_by_key[k]['auc'].append(results_i[k]['auc'])
        #        results_by_key[k]['ce'].append(results_i[k]['ce']);
        #    else:
        #        results_by_key[k]={'auc':[],'ce':[]};

        results_i=compute_metrics(scores.tolist(),gt.tolist(),data_test.data['table_ann']['trigger'])
        for k in results_i:
            if k in results_by_key:
                results_by_key[k]['auc'].append(results_i[k]['auc'])
                results_by_key[k]['ce'].append(results_i[k]['ce']);
            else:
                results_by_key[k]={'auc':[],'ce':[]};

        try:
            results_i=compute_metrics(scores.tolist(),gt.tolist(),data_test.data['table_ann']['trigger_subtype'])
        except:
            results_i={};

        for k in results_i:
            if k in results_by_key:
                results_by_key[k]['auc'].append(results_i[k]['auc'])
                results_by_key[k]['ce'].append(results_i[k]['ce']);
            else:
                results_by_key[k]={'auc':[],'ce':[]};

        auc.append(auc_i);
        ce.append(ce_i);
        cepre.append(ce_pre_i);
        session.log('Split %d, loss %.4f (%.4f), auc %.4f, time %f'%(split_id,ce_i,ce_pre_i,auc_i,time.time()-t0));

        ensemble.append({'net':net.cpu().state_dict(),'params':params_,'T':float(T.data.cpu())})

    auc=torch.Tensor(auc);
    ce=torch.Tensor(ce);
    cepre=torch.Tensor(cepre);

    if float(auc.mean())>best_loss_so_far:
        best_loss_so_far=float(auc.mean());
        torch.save(ensemble,session.file('model.pt'))

    session.log('AUC: %f + %f, CE: %f + %f, CEpre: %f + %f (%s (%d,%d,%d), epochs %d, batch %d, lr %f, decay %f)'%(auc.mean(),2*auc.std(),ce.mean(),2*ce.std(),cepre.mean(),2*cepre.std(),arch,nlayers,nlayers2,nh,epochs,batch,lr,decay));

    goal=-float(auc.mean()-2*auc.std());

    for k in results_by_key:
        auc=torch.Tensor(results_by_key[k]['auc']);
        ce=torch.Tensor(results_by_key[k]['ce']);
        session.log('\t KEY %s, AUC: %f + %f, CE: %f + %f'%(k,auc.mean(),2*auc.std(),ce.mean(),2*ce.std()));

    return goal;

#Get results from hyper parameter search
best=fmin(run_crossval,hp_config,algo=tpe.suggest,max_evals=params.budget)
#best=util.macro.obj(best);
params_=configure_pipeline(**best);
hyper_params_str=json.dumps(best);
session.log('Best hyperparam (%s)'%(hyper_params_str));

#Load extracted features
#fvs_0=torch.load('fvs.pt');
#fvs_1=torch.load('fvs_1.pt');
#fvs=db.union(db.Table.from_rows(fvs_0),db.Table.from_rows(fvs_1));
#fvs.add_index('model_id');

#Load labels
#label=[];
#for i in range(200):
#    fname='/work/projects/trojai-example/data/trojai-round0-dataset/id-%08d/ground_truth.csv'%i;
#    f=open(fname,'r');
#    for line in f:
#        line.rstrip('\n').rstrip('\r')
#        label.append(int(line));
#        break;
#
#    f.close();

#fvs['label']=label;
#data=db.DB({'table_ann':fvs});
#data.save('data.pt');


<filename>crystalmaths/imagetools.py
import skimage
import functools
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft2, fftshift, ifft
from skimage import img_as_float
from skimage.color import rgb2gray
from scipy.signal import get_window
from skimage import io
from skimage import transform

def setimage(google_shareable_link, output_filename, output_directory):
    """
    This function is similar to "show_image", but it loads image from a\\
     Google Drive instead of local directory.
    Inputs for the function are:
    (1) "Shareable link" from Google Drive provided as a string.
        Example: google_shareable_link =
        'https://drive.google.com/open?id=1cFi0rOqN8bcJ7H5fpfPAGS5Rem7TtiII'
        ***To get the link: go to the file in your Google Drive, right click,\\
        select "Get Shareable link".
    (2) Output file name including file extension provided as a string.
        Example: output_filename = 'Hexagonal_18.bmp'
    (3) Output Directory path provided as a string.
        Example: output_directory = '/Users/elenashoushpanova/Desktop/'
    Output is an array of image pixel values, image pixel resolution, as well\\
        as the image itself.

    Note: this function call for "file_download" and "show_image" functions.
    """
    # Call for a "file_download" function: ############ Will need a name \\
    # change to "crystalmath."
    dir_file = file_download(google_shareable_link, output_filename,
                             output_directory)

    # Call for a "show_image" function: ############ Will need a name change\\
    #  to "crystalmath."
    image = show_image(dir_file)
    return image, image.shape

def show_image(dir_file):
    """
    This function loads image from a local directory.
    Input is an image file path.
    Output is an array of image pixel values, as well as the image itself.
    """
    image = sio.imread(dir_file)
    sio.imshow(image)
    plt.axis('off')
    return image

class ImageHandler():
    """
    Class which takes as input an image filepath, turns it into an array, and
    processes the FFT. Methods include get_planes, get_ratio.

     If user chooses to apply window_type, window will be calculated based
     on argument provided. New windowed_image_array attribute created.
     The original array is still used for display purposes and identifying
     scalebar. But the image_fft will be of the windowed image instead of
     the original image.

     Todo:
    - create shifted points attribute which shifts the origin of the fft
    """
    def __init__(self, image_filepath, window_type=None, scale_ratio=None):
        self.image_filepath = image_filepath
        self.image_array = skimage.io.imread(self.image_filepath, as_gray=True)
        if window_type is not None:
            window = self.apply_window(window_type)
            self.windowed_image_array = self.image_array*window
            image_fft = np.abs(np.fft.fftshift((np.fft.fft2(
                self.windowed_image_array))))
            self.image_fft_array = np.log(image_fft)
        else:
            image_fft = np.abs(np.fft.fftshift((np.fft.fft2(
                self.image_array))))
            self.image_fft_array = np.log(image_fft)

    def show_image(self):
        fig, axes = plt.subplots(1, 2)
        axes[0].imshow(self.image_array, cmap='binary_r')
        axes[0].set_title('Loaded Image')
        axes[1].imshow(self.image_fft_array, cmap='binary_r')
        axes[1].set_title('Image FFT')
        plt.axis('off')

    def get_scale_ratio(self, scale_ratio=None):
        """
        This function can receive a pixel to nm ratio as its argument, which\
skips over defining the scalebar on the figure. If no ratio is provided,\
prompt is initiated. This will open up the image_array, and prompt user with\
instructions.
        """
        if scale_ratio is not None:
            # assert(type(scale_ratio)==float)
            self.scale_ratio = float(scale_ratio)
        else:
            prompt = 'Define the scalebar position. Pick end points using left\
mouse button. Right click once done. Middle mouse button removes most recent\
point.'
            fig, ax = plt.subplots()
            plt.setp(plt.gca(), autoscale_on=True)
            ax.imshow(self.image_array, cmap='binary_r')
            plt.title(prompt, wrap=True)
            points = []
            length_nm = float(input("Enter scalebar length in nanometers."))

            while len(points) < 2:
                points = np.asarray(plt.ginput(n=2, show_clicks=True,
                                    timeout=-1, mouse_add=1, mouse_stop=3,
                                    mouse_pop=2))
            length_pixel = np.abs(points[0][0]-points[1][0])
            print("length_nm = ", length_nm)
            print("length_pixel =", length_pixel)
            plt.close()
            self.scale_ratio = length_pixel/length_nm

    def get_planes(self):
        prompt = 'Pick two planes using left mouse button. Right click once\
 done. Middle mouse button removes most recent point.'
        message = "Press keyboard button to save points, mouseclick to restart"
        print(prompt)
        fig, ax = plt.subplots()
        plt.setp(plt.gca(), autoscale_on=True)
        ax.imshow(self.image_fft_array, cmap='binary_r')
        # zoom = False
        # while zoom is not True:
        while True:
            plt.title('Perform zoom if necessary, press down arrow key to\
 continue with plane selection', wrap=True)
            if plt.waitforbuttonpress(timeout=-1):
                # zoom = True
                break
        while True:
            plt.title(prompt, wrap=True)
            fig.canvas.draw()
            while True:
                points = plt.ginput(n=2, show_clicks=True, timeout=-1,
                                    mouse_add=1, mouse_stop=3, mouse_pop=2)
                break
            plt.title(message, wrap=True)
            fig.canvas.draw()
            print("Saved points = ", points)
            print(message)
            point_list = []
            if plt.waitforbuttonpress():
                for i in range(len(points)):
                    point_col = int(points[i][0])
                    point_row = int(points[i][1])
                    point_list.append([point_col, point_row])
                raw_point_coordinates = point_list
                # The center of the image is shifted to index 0,0, to ease
                # calculation of angles between points downstream.
                p1, p2 = raw_point_coordinates
                p1[0] = p1[0] - self.image_array.shape[0]/2
                p1[1] = p1[1] - self.image_array.shape[1]/2
                p2[0] = p2[0] - self.image_array.shape[0]/2
                p2[1] = p2[1] - self.image_array.shape[1]/2
                self.point_coordinates = [p1, p2]
                plt.close()
                break

    def apply_window(self, window_type):
        """
        Code taken from scikit-image window function. At time of writing\
        this the window function is only available in scikit-image development\
        build, and decision was made to not require that for this module.

        This is an optional method which can be called at the time of class\
        initialization in order to apply a specified window to the image array\
        before performing FFT. The purpose is to reduce image edge and\
        scalebar artifacts.
        """
        shape = self.image_array.shape
        ndim = len(shape)
        max_size = functools.reduce(max, shape)
        w = get_window(window_type, max_size, fftbins=False)
        w = np.reshape(w, (-1,)+(1,)*(ndim-1))

        # Create coords for warping following `ndimage.map_coordinates`
        # convention.
        L = [np.arange(s, dtype=np.float32) * (max_size / s) for s in shape]

        center = (max_size/2)-0.5
        dist = 0
        for g in np.meshgrid(*L, sparse=True, indexing='ij'):
            g -= center
            dist = dist + g*g
            dist = np.sqrt(dist)
            coords = np.zeros((ndim,)+dist.shape, dtype=np.float32)
            coords[0] = dist + center

        return transform.warp(w, coords, mode='constant')


<gh_stars>1-10
import os
import json
import numpy as np
from pathlib import Path
from scipy import ndimage
import SimpleITK as sitk

from report_guided_annotation import create_automatic_annotations_for_folder

def test_create_automatic_labels_for_folder():
    # define input and output folders
    input_dir = Path("tests/input/")
    output_dir = Path("tests/output/")

    # generate report-guided annotations
    create_automatic_annotations_for_folder(
        input_dir=input_dir,
        output_dir=output_dir,
        threshold='dynamic',
        skip_if_insufficient_lesions=True,
        num_workers=4,
    )

    # verify generated annotations exist
    assert os.path.exists(output_dir / "ProstateX-0000_07-07-2011-NA-MR prostaat kanker detectie WDSmc MCAPRODETW-05711.nii.gz")

    # read number of lesions there should be
    with open(input_dir / "num_lesions_to_retain_map.json") as fp:
        num_lesions_to_retain_map = json.load(fp)

    # for each case, check the number of lesions
    for subject_fn, num_lesions_to_retain in num_lesions_to_retain_map.items():
        # read automatic annotation
        automatic_annot = sitk.ReadImage(str(output_dir / subject_fn))
        automatic_annot = sitk.GetArrayFromImage(automatic_annot)
        _, num_blobs = ndimage.label(automatic_annot, structure=np.ones((3, 3, 3)))
        assert num_lesions_to_retain == num_blobs