import time
import os,binascii
from uuid import getnode as get_mac

from threading import Thread, Condition

from receive_thread import *
from transmit_thread import *

from constants import *

import re

# use last 4 hex digits of MAC address as unique car id
DEVICE_ID = hex(get_mac())[-5:-1]
MESSAGE_FORMAT = re.compile('^([a-fA-F0-9]{4}):([0-3]{1}):([0-3]{1}):([0-9]{3})$')

last_message = ""

class V2VModule(Thread):

	def __init__(self, ready_callback, debug_mode=True):
		super(V2VModule, self).__init__()

		self.ready_callback = ready_callback
		self.debug_mode = debug_mode

		# transmitter and receiver threads
		#receiver callback : _on_message_received
		self.condition = Condition() # thread control variable (ensures alternating threads)
		self.receiver = ReceiveThread(self.condition, self._on_message_received)
		self.transmitter = TransmitThread(None, self.condition, None)

		# car state, set externally ONLY by main car controller
		# IDLE, STOPPED, IN_TRANSIT, CLEARED
		self.state = IDLE

		# when we are ready to go through intersection
		self.ready = False

		# data for determining queue/departure order
		self.vehicles = {}
		self.arrival_time = None
		self.direction = None

		# thread run condition
		self.running = True

	def run(self):
		self._set_transmitter_state(IDLE)
		self.transmitter.start()
		self.receiver.start()
		self.debug_print("V2V Module started.")
		while self.running:

			# do nothing until we get to an intersection
			# during this time vehicle queue is constantly being updated as _on_message_received is triggered
			time.sleep(0.05)

			# transmit STOPPED signal and wait for our turn to transit
			if self.state == STOPPED:
				self.debug_print("Stopped. Waiting for turn to transit.")
				self._set_transmitter_state(STOPPED)

				time.sleep(MINIMUM_WAIT_PERIOD) # minimum time we have to wait
				self.debug_print("Minimum wait time (MINIMUM_WAIT_PERIOD = %.1f) elapsed." % (MINIMUM_WAIT_PERIOD))

				# wait for our turn in queue to go, either when its our turn or if there is no one else
				while not ( self.ready or len(self.vehicles.keys()) == 0 ):
					time.sleep(0.1)

				self.debug_print("ready flag set. broadcasting IN_TRANSIT signal.")
				self.ready_callback()
				# our turn to go. set TRANSIT signal
				self._set_transmitter_state(IN_TRANSIT)

				# wait for car controller to receive ready signal and drive through intersection
				while self.state != IN_TRANSIT:
					time.sleep(0.1)
				self.debug_print("Entered intersection...")

				# car is moving. wait until we have left intersection
				while self.state != CLEARED:
					time.sleep(0.1)
				self.debug_print("Cleared intersection. Broadcasting CLEARED signal.")
				# through intersection. broadcast CLEARED signal for 3 seconds.
				self._set_transmitter_state(CLEARED)
				time.sleep(3)

				self.debug_print("Returning to idle state.")
				# go back to idle state until we arrive at intersection again
				self._set_idle()

	# bulk of our calculations go here
	# determine vehicle transit order based on arrival time and lane precedence
	#
	# example of data structure :
	#
	# self.vehicles = {
	# 	"id_x" : {
	# 		state: STOPPED,
	# 		direction: SOUTH,
	# 		speed: 555
	# 	},
	# 	...
	# }
	def _update_vehicles(self, vehicle_id, vehicle_data):

		#self.debug_print("new vehicle data received (%s): %s " % ( str(vehicle_id), str(vehicle_data) ) )
		#self.debug_print("# of vehicles in vehicle data ")
		#self.debug_print("_update_vehicles called. # of vehicles recorded: %d" % (len(self.vehicles.keys())) )
		self.ready = False

		state = vehicle_data["state"]

		# vehicle was already recorded. update its state
		if vehicle_id in self.vehicles.keys():
			if state == CLEARED: #remove the vehicle from our record
				self.debug_print("cleared state read for %s. removing." % (vehicle_id))
				del self.vehicles[vehicle_id]
			elif state != self.vehicles[vehicle_id]["state"]:
				self.debug_print("updating state for %s" % (vehicle_id))
				self.vehicles[vehicle_id]["state"] = vehicle_data["state"]

		# new vehicle. record time of arrival
		elif state != CLEARED:
			self.debug_print("new vehicle arrived and stopped: %s" % (vehicle_id))
			self.vehicles[vehicle_id] = vehicle_data
			self.vehicles[vehicle_id]["timestamp"] = time.time()

		if self.arrival_time is not None:

			# TODO move to seperate function (like in MID, MIS)
			#get the earliest arrival time
			earliest_arrival = min([self.vehicles[v]["timestamp"] for v in self.vehicles])
			#earliest_arrival = min([v["timestamp"] for v in self.vehicles])

			# check if there is already a car in the intersection
			transit_vehicle = None
			for v in self.vehicles:
				if self.vehicles[v]["state"] == IN_TRANSIT:
					transit_vehicle = self.vehicles[v]

			if transit_vehicle:
				self.debug_print("Another vehicle currently in transit.")
				if (transit_vehicle["direction"] in [NORTH,SOUTH] and self.direction in [NORTH, SOUTH]) \
						or (transit_vehicle["direction"] in [EAST,WEST] and self.direction in [EAST, WEST]):
					self.debug_print("In-transit vehicle in parallel lane. ready.")
					self.ready =  True

			# check for arrival time precedence
			elif (self.arrival_time + BUFFER_PERIOD) < earliest_arrival :
				self.debug_print("arrival time + Buffer period is less than next-earliest arrival. going.")
				self.ready = True

			# check for lane precedence (NORTH/SOUTH gets to go first)
			elif abs(self.arrival_time - earliest_arrival) < BUFFER_PERIOD:
				self.debug_print("ambiguous arrival times")
				if self.direction in [NORTH,SOUTH]:
					self.debug_print("we are in priority lane. going.")
					self.ready = True

	# Receiver callback, triggered when a NEW message is received
	# update our vehicle data
	def _on_message_received(self, msg):
		string = ''.join(chr(e) for e in msg)
		params = self._parse_message(str(string))

		if params:
			#self.debug_print("new message received: %s" % (string))
			vehicle_id = str(params[0])
			vehicle_data = {
					"state": int(params[1]),
					"direction": int(params[2]),
					"speed": int(params[3])
			}

			self._update_vehicles(vehicle_id, vehicle_data)

	# validate format, extract vehicle data and remove self-messages
	def _parse_message(self, msg):
		# regex match
		match = MESSAGE_FORMAT.match(msg)
		if match and str(match.groups()[0]) != DEVICE_ID:
			return match.groups()
		else:
			return False

	# set the message being broadcasted continuously by transmitter
	def _set_transmitter_state(self, state):
		if state == IDLE:
			self.transmitter.set_enabled(False)
		else:
			# e.g:  ab12:1:3:777
			message = str(DEVICE_ID) + ":" + str(state) + ":" + str(self.direction) + ":" + str(DUMMY_SPEED)
			self.transmitter.set_message(message)
			self.transmitter.set_enabled(True)

	# reset after leaving intersection
	def _set_idle(self):
		self.direction = None
		self.arrival_time = None
		self.ready = False
		self.state = IDLE
		self._set_transmitter_state(IDLE)

	"""
	public methods called by main vehicle controller
	"""

	# check if we are allowed to drive into intersection
	def get_transit_permission(self):
		return self.ready

	# 1 arrive at intersection
	def set_stopped(self, direction):
		self.debug_print('Stopped state received from main controller')
		self.direction = direction
		self.arrival_time = time.time()
		self.state = STOPPED

	# 2 driving through intersection after receiving permission
	def set_in_transit(self):
		self.debug_print('In transit state received from main controller')
		self.state = IN_TRANSIT

	# 3 left intersection
	def set_cleared(self):
		self.debug_print('Cleared state received from main controller')
		self.state = CLEARED

	# debugging
	def debug_print(self, string):
		global last_message
		if self.debug_mode and string != last_message:
			last_message = string
			print( str(DEVICE_ID) + ": " + string)


from datetime import datetime
import numpy as np
import torch
from matplotlib import pyplot as plt
from torch import nn

from _01_code._06_fcn_best_practice.c_trainer import EarlyStopping
from _01_code._99_common_utils.utils import strfdelta

class AutoencoderTrainer:
  def __init__(
    self, project_name, model, optimizer, train_data_loader, validation_data_loader, transforms,
    run_time_str, wandb, device, checkpoint_file_path, test_dataset, test_transforms, denoising=True
  ):
    self.project_name = project_name
    self.model = model
    self.optimizer = optimizer
    self.train_data_loader = train_data_loader
    self.validation_data_loader = validation_data_loader
    self.transforms = transforms
    self.run_time_str = run_time_str
    self.wandb = wandb
    self.device = device
    self.checkpoint_file_path = checkpoint_file_path

    self.test_dataset = test_dataset
    self.test_transforms = test_transforms
    self.denoising = denoising

    # Use a built-in loss function
    self.loss_fn = nn.MSELoss()

  def add_noise(self, inputs, noise_factor=0.1):
    noisy = inputs + torch.randn(inputs.size()) * noise_factor
    noisy = torch.clip(noisy, 0., 1.)
    return noisy

  def do_train(self):
    self.model.train()

    loss_train = 0.0
    num_trains = 0

    # Iterate the dataloader (we do not need the label values, this is unsupervised learning)
    for train_batch in self.train_data_loader:
      # with "_" we just ignore the target labels
      input_train, _ = train_batch

      if self.denoising is True:
        input_train = self.add_noise(input_train)

      input_train = input_train.to(device=self.device)

      if self.transforms:
        input_train = self.transforms(input_train)

      decoded_input_train = self.model(input_train)

      loss = self.loss_fn(decoded_input_train, input_train)

      loss_train += loss.item()

      num_trains += 1

      self.optimizer.zero_grad()
      loss.backward()
      self.optimizer.step()

    train_loss = loss_train / num_trains

    return train_loss

  def do_validation(self):
    self.model.eval()

    loss_validation = 0.0
    num_validations = 0

    with torch.no_grad():
      for validation_batch in self.validation_data_loader:
        input_validation, _ = validation_batch

        if self.denoising is True:
          input_validation = self.add_noise(input_validation)

        input_validation = input_validation.to(device=self.device)

        if self.transforms:
          input_validation = self.transforms(input_validation)

        decoded_input_validation = self.model(input_validation)

        loss_validation += self.loss_fn(decoded_input_validation, input_validation).item()

        num_validations += 1

    validation_loss = loss_validation / num_validations

    return validation_loss

  def plot_denoising_autoencoders_outputs(self, n=10, noise_factor=0.1):
    self.model.eval()

    plt.figure(figsize=(16, 4.5))
    targets = self.test_dataset.targets.numpy()
    t_idx = {i: np.where(targets == i)[0][0] for i in range(n)}
    for i in range(n):
      ax = plt.subplot(3, n, i + 1)
      img = self.test_dataset.data[t_idx[i]].unsqueeze(0).unsqueeze(0)

      if self.denoising is True:
        image_noisy = self.add_noise(img, noise_factor)
        image_noisy = image_noisy.to(self.device)
      else:
        img = img.type(torch.float)
        img = img.to(self.device)

      with torch.no_grad():
        decoded_img = self.model(image_noisy if self.denoising is True else img)

      plt.imshow(img.cpu().squeeze().numpy(), cmap='gist_gray')
      ax.get_xaxis().set_visible(False)
      ax.get_yaxis().set_visible(False)
      if i == n // 2:
        ax.set_title('Original images')

      if self.denoising is True:
        ax = plt.subplot(3, n, i + 1 + n)
        plt.imshow(image_noisy.cpu().squeeze().numpy(), cmap='gist_gray')
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        if i == n // 2:
          ax.set_title('Corrupted images')

      ax = plt.subplot(3, n, i + 1 + n + n)
      plt.imshow(decoded_img.cpu().squeeze().numpy(), cmap='gist_gray')
      ax.get_xaxis().set_visible(False)
      ax.get_yaxis().set_visible(False)
      if i == n // 2:
        ax.set_title('Reconstructed images')

    plt.subplots_adjust(left=0.1, bottom=0.1, right=0.7, top=0.9, wspace=0.3, hspace=0.3)
    plt.show(block=False)

  def train_loop(self):
    early_stopping = EarlyStopping(
      patience=self.wandb.config.early_stop_patience,
      delta=self.wandb.config.early_stop_delta,
      project_name=self.project_name,
      checkpoint_file_path=self.checkpoint_file_path,
      run_time_str=self.run_time_str
    )
    n_epochs = self.wandb.config.epochs
    training_start_time = datetime.now()

    for epoch in range(1, n_epochs + 1):
      train_loss = self.do_train()

      if epoch == 1 or epoch % self.wandb.config.validation_intervals == 0:
        validation_loss = self.do_validation()

        elapsed_time = datetime.now() - training_start_time
        epoch_per_second = 1000 * epoch / elapsed_time.microseconds

        message, early_stop = early_stopping.check_and_save(validation_loss, self.model)

        print(
          f"[Epoch {epoch:>3}] "
          f"T_loss: {train_loss:7.5f}, "
          f"V_loss: {validation_loss:7.5f}, "
          f"{message} | "
          f"T_time: {strfdelta(elapsed_time, '%H:%M:%S')}, "
          f"T_speed: {epoch_per_second:4.3f}"
        )

        self.wandb.log({
          "Epoch": epoch,
          "Training loss": train_loss,
          "Validation loss": validation_loss,
          "Training speed (epochs/sec.)": epoch_per_second,
        })

        self.plot_denoising_autoencoders_outputs(n=10, noise_factor=0.3)

        if early_stop:
          break

    elapsed_time = datetime.now() - training_start_time
    print(f"Final training time: {strfdelta(elapsed_time, '%H:%M:%S')}")


#!/usr/bin/env python

'''
Created 09 December 2018
ptashkir@mskcc.org
wrapper for cfdna copy number pipeline
#######################################
to do:
take single bed file and calculate GC of intervals at run time
incoporate off target intervals, if needed

'''

import argparse
import os
import re
import sys
import time
import stat
import csv
from subprocess import Popen
import shlex
import pysam
import pandas as pd
from joblib import Parallel, delayed

def main():
    parser = argparse.ArgumentParser(prog='cfdna_scna.py', description='cfdna copy number pipeline wrapper', usage='%(prog)s [options]')
    parser.add_argument("-t", "--tumorManifest", action="store", dest="tumorManifest", required=True, metavar='tumors.txt', help="Full path to the tumor sample manifest, tab serparated BAM path, patient sex.")
    parser.add_argument("-n", "--normalManifest", action="store", dest="normalManifest", required=True, metavar='normals.txt', help="Full path to the normal sample manifest, tab serparated BAM path, patient sex.")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose", default=True, help="make lots of noise [default]")
    parser.add_argument("-tr", "--threads", action="store", dest="threads", required=False, metavar='8', default='8', help="Number of Threads to be used to generate coverage metrics")
    parser.add_argument("-b", "--bedTargets", action="store", dest="TARGETS", required=True, metavar='/somepath/ACCESS_targets_coverage.bed', help="Full Path to BED file of panel targets")
    #parser.add_argument("-ta", "--targetAnnotations", action="store", dest="ANNOTATIONS", required=True, metavar='/somepath/ACCESS_targets_coverage.txt', help="Full Path to text file of target annotations. Columns = (Chrom, Start, End, Target, GC_150bp, GeneExon,Cyt,Interval)")
    parser.add_argument("-g", "--genomeReference", action="store", dest="GENOME", required=True, metavar='/somepath/Homo_Sapeins_hg19.fasta', help="Full Path to the reference fasta file.")
    #parser.add_argument("-r", "--RPATH", action="store", dest="R", required=True, metavar='/somepath/R', help="Path to R executable.")
    #parser.add_argument("-q", "--queue", action="store", dest="queue", required=True, metavar='test.q or clin2.q', help="Name of the SGE queue")
    parser.add_argument("-id", "--runID", action="store", dest="runID", required=True, metavar='ACCESSv1-VAL-20180001', help="name of runID.")
    #parser.add_argument("-o", "--outDir", action="store", dest="outdir", required=True, metavar='/somepath/output', help="Full Path to the output dir.")
    #parser.add_argument("-l", "--loess", action="store", dest="loess", required=True, metavar='/somepath/loessnormalize_nomapq_cfdna.R', help="Full Path to the loess normalization R script.")
    #parser.add_argument("-cn", "--copynumber", action="store", dest="cnAnalysis", required=True, metavar='/somepath/copynumber_tm.batchdiff_cfdna.R', help="Full Path to the copy number R script.")
    #parser.add_argument("-qsub", "--qsubPath", action="store", dest="qsub", required=False, metavar='/somepath/qsub', help="Full Path to the qsub executables of SGE.")
    #parser.add_argument("-bsub", "--bsubPath", action="store", dest="bsub", required=False, metavar='/somepath/bsub', help="Full Path to the bsub executables of LSF.")
    #parser.add_argument("-gatk", "--GATK", action="store", dest="GATK", required=False, metavar='/somepath/GATK', help="Full Path to the GATK.")
    #parser.add_argument("-j", "--javaPATH", action="store", dest="JAVA", required=False, metavar='/somepath/java', help="Path to java executable.")

    print("Running the cfdna copy number pipeline.")
    sys.stdout.flush()
    args = parser.parse_args()
    threads= int(args.threads)

    with Parallel(n_jobs=threads,verbose=1) as parallel:
        (patientSex)= ProcessArgs(args)

        (tumorCovFile) = RunBedCov(args,'tumor_bams.list', 'tumors',parallel)
        (normalCovFile) = RunBedCov(args,'normal_bams.list', 'normals',parallel)

        #(normalLoessFile) = RunLoessNormalization(args, normalCovFile, 'normal')
        #(tumorLoessFile) = RunLoessNormalization(args, tumorCovFile, 'tumor')

        #RunTumorCN(args, normalLoessFile, tumorLoessFile)

def ProcessArgs(args):

    tumorBamFiles = []
    normalBamFiles = []
    patientSex={}
    if(args.verbose):
        print("Going to see how many bam files are there for us to run the analysis.")
        sys.stdout.flush()

    with open(args.tumorManifest, 'r') as tfile:
        tlist= open('tumor_bams.list', 'w')
        for line in tfile:
            data = line.rstrip('\n')
            bfile = re.split('\t', data)[0]
            sex = re.split('\t', data)[1]
            tumorBamFiles.append(bfile)
            patientSex.update({bfile:sex})
            tlist.write(bfile+"\n")
        tlist.close()

    with open(args.normalManifest, 'r') as nfile:
        nlist= open('normal_bams.list', 'w')
        for line in nfile:
            data = line.rstrip('\n')
            bfile = re.split('\t', data)[0]
            normalBamFiles.append(bfile)
            sex = re.split('\t', data)[1]
            patientSex.update({bfile:sex})
            nlist.write(bfile+"\n")
        nlist.close()

    if(args.verbose):
        print("Going to run analysis for", tumorBamFiles.__len__(), " tumor BAM files using ", normalBamFiles.__len__()," normal BAM files")

    return(patientSex)

def RunBedCov(args,bamlist, runType, parallel):
    if(args.verbose):
        print("generating coverage metrics for " +runType+ " BAMS using pysam bedcov")
        sys.stdout.flush()

    covFile = args.runID + "_" +runType+ "_targets_nomapq.covg_interval_summary"

    if os.path.exists(covFile):
        print("Coverage analysis file " + covFile + " already exists, will not generate coverage metrics")
        return(covFile)

    bfile = open(bamlist)
    bed = args.TARGETS

    outname = args.runID + "_" +runType+ "_targets_nomapq.covg"
    cov_args = []
    for line in bfile:
        for bam in line.split():
            cov_args.append((bam, bed))

    results = parallel(map(delayed(parallelCov), cov_args))
    results= {k: v for x in results for k, v in x.items()}

    df= pd.DataFrame.from_records(results, index='Target')
    df.to_csv(covFile, sep='\t')

    if os.path.isfile(covFile):
        print("Completed generation of coverage metrics for " +runType+ " BAMS using pysam bedcov")
        return(covFile)
    else:
        print("Coverage file " + covFile + "does note exist, something went wrong here")

def parallelCov(cov_args):
    (bam,bed) = cov_args
    sampleCov=[]
    output={}
    mq = 0 #can param min mapping quality if desired
    id= os.path.basename(bam)
    id = id.split('_')[0]
    print("Generating coverage metrics for: " + id)
    sys.stdout.flush()
    id = id+"_mean_cvg"
    cmd = [bed, bam]
    cmd.extend(['-Q',bytes(mq)])

    bstring = pysam.bedcov(*cmd, split_lines=False)
    lines = bstring.splitlines()
    targets=[]
    for line in lines:
        fields = line.split('\t')
        chr = fields[0]
        start = int(fields[1])
        end= int(fields[2])
        target = fields[3]
        coverage = int(fields[4])
        intlen = float(end - start)
        meancov = str(coverage / intlen)
        target = chr +":"+str(start)+"-"+str(end)
        targets.append(target)
        sampleCov.append(meancov)
    output.update({id:sampleCov})
    output.update({'Target':targets})

    return(output)

def RunLoessNormalization(args, covFile, runType):
    loessFile = args.runID + "_" + runType + "_ALL_intervalnomapqcoverage_loess.txt"

    if(args.verbose):
        print("running loess normalization")
        sys.stdout.flush()

    cmd = args.R + " --slave --vanilla --args " + args.runID+ " " + args.outdir + " " + args.ANNOTATIONS + " " + covFile + " " + runType + "< " + args.loess

    cl_cmd =''
    mem = 8
    maxmem = int(mem)+8
    if(args.qsub):
        cl_cmd = args.qsub + " -q " + args.queue + " -N " + "loessNorm" + " -V -l h_vmem=8G,virtual_free=8G -pe smp 1 -wd " + args.outdir + " -sync y " + " -b y " + cmd
        print("CLUSTER_CMD loess==>", cl_cmd , "\n")

    else:
        cl_cmd= args.bsub + " -q " + args.queue + " -J " + "loessNorm" + " -o loessNorm.%J.out -e loessNorm.%J.err" + " -We 24:00 -R \"rusage[mem=" + str(mem) + "]\" -M " + str(maxmem) + " -n 1 " + " -cwd " + args.outdir + ' -K  "' + cmd + '"'
        print("CLUSTER_CMD loess==>", cl_cmd , "\n")
    cl_args = shlex.split(cl_cmd)
    proc = Popen(cl_args)
    proc.wait()
    retcode = proc.returncode
    if(retcode == 0):
        if(args.verbose):
            print("Finished Running loess normalization for " + runType + " BAM files")
    else:
        if(args.verbose):
            print("loess normalization is either still running or it errored out with return code", retcode,"\n")

    if os.path.isfile(loessFile):
        return(loessFile)
    else:
        print("Loess normalization file " + loessFile + " does not exist, something went wrong here")

def RunTumorCN(args,normalLoessFile, tumorLoessFile):
    if(args.verbose):
        print("calculating log ratio, performing segmentation and plot fest!")
        sys.stdout.flush()

    cmd = args.R + " --slave --vanilla --args " + args.runID+ " " + normalLoessFile + " " + args.ANNOTATIONS + " " + tumorLoessFile + " MIN "+ "< " + args.cnAnalysis

    cl_cmd =''
    mem = 8
    maxmem = int(mem)+8
    if(args.qsub):
        cl_cmd = args.qsub + " -q " + args.queue + " -N " + "runCN" + " -V -l h_vmem=8G,virtual_free=8G -pe smp 1 -wd " + args.outdir + " -sync y " + "-b y " + cmd
        print("CLUSTER_CMD CN==>", cl_cmd , "\n")

    else:
        cl_cmd= args.bsub + " -q " + args.queue + " -J " + "runCN" + " -o runCN.%J.out -e runCN.%J.err"+ " -We 24:00 -R \"rusage[mem=" + str(mem) + "]\" -M " + str(maxmem) + " -n 1 " + " -cwd " + args.outdir + ' -K "' + cmd + '"'
        print("CLUSTER_CMD CN==>", cl_cmd , "\n")
    cl_args = shlex.split(cl_cmd)
    proc = Popen(cl_args)
    proc.wait()
    retcode = proc.returncode
    if(retcode == 0):
        if(args.verbose):
            print("Finished Running segmentation and plotting")
            print("Done running Copy Number pipeline.")
    else:
        if(args.verbose):
            print("copy number analysis is either still running or it errored out with return code", retcode,"\n")

def RunCoverage(args): #this method is now depricated
    if(args.verbose):
        print("generating coverage metrics for BAMS")

    #cmd_t = args.JAVA + " -Djava.io.tmpdir=/dmp/analysis/SCRATCH -Xmx4g -jar " + args.GATK + " -T DepthOfCoverage -R " + args.GENOME + " -I tumor_bams.list"  + " -o " + args.runID + "_tumors_targets_nomapq.covg" + " -L " + args.TARGETS + " -rf BadCigar -mmq 0 -mbq 20 -omitLocusTable -omitSampleSummary -omitBaseOutput --allow_potentially_misencoded_quality_scores"
    cmd_t = args.JAVA + " -Djava.io.tmpdir=/scratch -Xmx4g -jar " + args.GATK + " -T DepthOfCoverage -R " + args.GENOME + " -I tumor_bams.list"  + " -o " + args.runID + "_tumors_targets_nomapq.covg" + " -L " + args.TARGETS + " -rf BadCigar -mmq 0 -mbq 20 -omitLocusTable -omitSampleSummary -omitBaseOutput --allow_potentially_misencoded_quality_scores"
    #cmd_n = args.JAVA + " -Djava.io.tmpdir=/dmp/analysis/SCRATCH -Xmx4g -jar " + args.GATK + " -T DepthOfCoverage -R " + args.GENOME + " -I normal_bams.list"  + " -o " + args.runID + "_normals_targets_nomapq.covg" + " -L " + args.TARGETS + " -rf BadCigar -mmq 0 -mbq 20 -omitLocusTable -omitSampleSummary -omitBaseOutput --allow_potentially_misencoded_quality_scores"
    cmd_n = args.JAVA + " -Djava.io.tmpdir=/scratch -Xmx4g -jar " + args.GATK + " -T DepthOfCoverage -R " + args.GENOME + " -I normal_bams.list"  + " -o " + args.runID + "_normals_targets_nomapq.covg" + " -L " + args.TARGETS + " -rf BadCigar -mmq 0 -mbq 20 -omitLocusTable -omitSampleSummary -omitBaseOutput --allow_potentially_misencoded_quality_scores"

    cl_cmd_t = ''
    cl_cmd_n = ''

    mem = int(args.threads) * 8
    maxmem = int(mem)+8
    if(args.qsub):
        cl_cmd_t = args.qsub + " -q " + args.queue + " -N " + "DOC_tumors" + " -V -l h_vmem=8G,virtual_free=8G -pe smp " + args.threads + " -wd " + args.outdir + " -sync y " + " -b y " + cmd_t
        cl_cmd_n = args.qsub + " -q " + args.queue + " -N " + "DOC_normals" + " -V -l h_vmem=8G,virtual_free=8G -pe smp " + args.threads + " -wd " + args.outdir + " -sync y " + " -b y " + cmd_n
    else:
        cl_cmd_t = args.bsub + " -q " + args.queue + " -J " + "DOC_tumors" + " -We 24:00 -R \"rusage[mem=" + str(mem) + "]\" -M " + str(maxmem) + " -n " + args.threads + " -cwd " + args.outdir + " -K " + cmd_t
        cl_cmd_n = args.bsub + " -q " + args.queue + " -J " + "DOC_tumors" + " -We 24:00 -R \"rusage[mem=" + str(mem) + "]\" -M " + str(maxmem) + " -n " + args.threads + " -cwd " + args.outdir + " -K " + cmd_n

    tumorCovFile = args.outdir + "/"+ args.runID + "_tumors_targets_nomapq.covg.sample_interval_summary"
    normalCovFile = args.outdir + "/"+ args.runID + "_normals_targets_nomapq.covg.sample_interval_summary"

    if os.path.exists(tumorCovFile):
        print("Coverage analysis for tumor bams already exist, will not run")

    else:
        print("CLUSTER_CMD tumors==>", cl_cmd_t , "\n")
        cl_args_t = shlex.split(cl_cmd_t)
        proct = Popen(cl_args_t)
        proct.wait()
        retcodet = proct.returncode
        if(retcodet == 0):
            if(args.verbose):
                print("Finished Running coverage analysis for tumor bams ")
        else:
            if(args.verbose):
                print("coverage analysis is either still running or it errored out with return code", retcode,"\n")

    if os.path.exists(normalCovFile):
        print("Coverage analysis for normal bams already exist, will not run")
    else:
        print("CLUSTER_CMD normals==>", cl_cmd_n , "\n")
        cl_args_n = shlex.split(cl_cmd_n)
        procn = Popen(cl_args_n)
        procn.wait()
        retcoden = procn.returncode
        if(retcoden == 0):
            if(args.verbose):
                print("Finished Running coverage analysis for normal bams ")
        else:
            if(args.verbose):
                print("coverage analysis is either still running or it errored out with return code", retcode,"\n")

    files = [normalCovFile,tumorCovFile]
    if os.path.isfile(normalCovFile and tumorCovFile):
        return(tumorCovFile, normalCovFile)
    else:
        print("Tumor and Normal BAM coverage files do not exist, something went wrong here")

if __name__ == "__main__":
    start_time = time.time()
    main()
    end_time = time.time()
    print("Elapsed time was %g seconds" % (end_time - start_time))


from django import forms

from . import models

class TinyMCEForm(forms.ModelForm):
    class Media:
        css = {
            'all': (
                '/static/css/autocomplete.css',
                'https://cdnjs.cloudflare.com/ajax/libs/at.js/1.5.2/css/jquery.atwho.min.css',
                )
            }
        js = (
            'https://cdnjs.cloudflare.com/ajax/libs/at.js/1.5.2/js/jquery.atwho.min.js',
            'https://cdnjs.cloudflare.com/ajax/libs/Caret.js/0.3.1/jquery.caret.min.js',
            '/static/js/tinymce/tinymce.min.js',
            )

class MonsterForm(TinyMCEForm):
    class Meta:
        model = models.Monster
        fields = [
            'name',
            'alignment',
            'size',
            'languages',
            'strength',
            'dexterity',
            'constitution',
            'intelligence',
            'wisdom',
            'charisma',
            'armor_class',
            'hit_points',
            'speed',
            'saving_throws',
            'skills',
            'creature_type',
            'damage_vulnerabilities',
            'damage_immunities',
            'damage_resistances',
            'condition_immunities',
            'senses',
            'challenge_rating',
            'content',
        ]

class NPCForm(TinyMCEForm):
    class Meta:
        model = models.NPC
        fields = [
            'name',
            'alignment',
            'size',
            'languages',
            'strength',
            'dexterity',
            'constitution',
            'intelligence',
            'wisdom',
            'charisma',
            'race',
            'armor_class',
            'hit_points',
            'speed',
            'saving_throws',
            'skills',
            'npc_class',
            'age',
            'height',
            'weight',
            'creature_type',
            'damage_vulnerabilities',
            'damage_immunities',
            'damage_resistances',
            'condition_immunities',
            'senses',
            'challenge_rating',
            'content',
        ]

class PlayerForm(TinyMCEForm):
    class Meta:
        model = models.Player
        fields = [
            'player_name',
            'character_name',
            'level',
            'xp',
            'race',
            'character_class',
            'background',
            'alignment',
            'armor_class',
            'initiative',
            'speed',
            'proficiency_bonus',
            'hit_points',
            'strength',
            'dexterity',
            'constitution',
            'intelligence',
            'wisdom',
            'charisma',
            'saving_throws',
            'skills',
            'age',
            'height',
            'weight',
            'languages',
            'personality',
            'ideals',
            'bonds',
            'flaws',
            'feats',
            'attacks',
            'spells',
            'proficiencies',
            'equipment',
            'notes',
        ]

class CopyMonsterForm(forms.ModelForm):
    class Meta:
        model = models.Monster
        fields = ['name']

class CopyNPCForm(forms.ModelForm):
    class Meta:
        model = models.NPC
        fields = ['name']

class CopyPlayerForm(forms.ModelForm):
    class Meta:
        model = models.Player
        fields = ['player_name']

class ImportMonsterForm(forms.ModelForm):
    class Meta:
        model = models.Monster
        fields = [
            'name',
        ]

class ImportNPCForm(forms.ModelForm):
    class Meta:
        model = models.NPC
        fields = [
            'name',
        ]

class SRDMonsterForm(forms.ModelForm):
    class Meta:
        model = models.Monster
        fields = ['name']

class SRDNPCForm(forms.ModelForm):
    class Meta:
        model = models.NPC
        fields = ['name']

MonsterFormSet = forms.modelformset_factory(
    models.Monster,
    form=MonsterForm,
    extra=0,
)

NPCFormSet = forms.modelformset_factory(
    models.NPC,
    form=NPCForm,
    extra=0,
)

class MonsterPublishForm(TinyMCEForm):
    class Meta:
        model = models.Monster
        fields = [
            'tavern_description',
        ]

class NPCPublishForm(TinyMCEForm):
    class Meta:
        model = models.NPC
        fields = [
            'tavern_description',
        ]

class PlayerPublishForm(TinyMCEForm):
    class Meta:
        model = models.Player
        fields = [
            'tavern_description',
        ]


# Homework 2 program to determine accuracy for clustering.

import numpy
import sys
from munkres import Munkres

# Paths for the input files.
data_label_path = './data_label.txt'
clustering_result_path = './clustering_result.txt'
output_file_path = './clustering_accuracy.txt'

# Lists to hold contents of input files.
data_label = []
clustering_result = []

# Sets to verify count of labels and count of clusters match.
data_label_set = set()
clustering_result_set = set()

# Read data label file.
try:
    with open(data_label_path, 'r') as data_label_input:
        for current_label in data_label_input:
            data_label.append(int(current_label.split(' ')[0]))
            data_label_set.add(int(current_label.split(' ')[0]))
except:
    error_message = 'File + ' + data_label_path + ' not found.'
    exit(error_message)
data_label_input.close()

# Read clustering results file
try:
    with open(clustering_result_path, 'r') as clustering_result_input:
        for current_result in clustering_result_input:
            clustering_result.append(int(current_result.rstrip('\n')))
            clustering_result_set.add(int(current_result.rstrip('\n')))
except:
    error_message = 'File + ' + clustering_result_path + ' not found.'
    exit(error_message)
clustering_result_input.close()

explanation = """
Nick Petty
CAP6776 Homework 3

Accuracy = max (sum{Ck, Lm} T(Ck, Lm)) /n
Where n = number of data points, Ck denotes kth cluster, Lm is the mth class, and T(Ck, Lm) is the number of data points that belong to class m and are assigned to cluster k.  Because k and m refer to the number of classes in the data set, they must be equal.

Finding the maximum sum of T(Ck, Lm), when given T in matrix form, is a variation on the Assignment Problem.  Solving this problem via brute force will give a run time of m!, which is excessive.  Instead, the Munkres algorithm (or the Hungarian Algorithm) is used.  This gives a run time of m^3, which is much better.

To use the Munkres algorithm, the Python munkres package has been imported.
Package documentation: https://pypi.python.org/pypi/munkres/
The matrix from T(Ck, Lm) does not immediately work with this algorithm, as it is designed to find the minimum sum.  A second matrix, composed of a very large number minus each cell value in the T matrix must also be created.  This larger-valued matrix, when given to the algorithm, will then find the maximum sum.\n"""

# Set variables for calculation:

# Verify that count of labels matches count of clusters.  Find the number of classes m and clusters k.
if len(data_label_set) == len(clustering_result_set):
    m = len(data_label_set)
else:
# Write the output file and end the program.
    with open(output_file_path, 'w') as output_file:
        output_file.write(explanation)
        output_file.write("The number of data labels is not the same as the number of clusters.  Clustering accuracy cannot be calculated.")
    output_file.close()
    quit()

# Sum of T(Ck, Lm) values to maximize
max_sum = 0

# Number of data points.
n = len(data_label)

# Matrix to hold T(Ck, Lm)).
t_matrix = numpy.zeros((m, m), dtype=numpy.int)

# Matrix to hold large-valued t_matrix.
t_matrix_large = numpy.zeros((m, m), dtype=numpy.int)

# Generate T(Ck, Lm) and the large version as a copy.
for index in range(0, n):
    t_matrix[clustering_result[index] - 1][data_label[index] - 1] += 1
    t_matrix_large[clustering_result[index] - 1][data_label[index] - 1] += 1

# Fill the large version of T(Ck, Lm) with large values.
for i in range(0, m):
    for j in range(0, m):
        t_matrix_large[i][j] = sys.maxsize - t_matrix_large[i][j]

# Create a munkres object.
munkres_object = Munkres()

# Use the Munkres algorithm to find the location of cells that make the maximum sum.
max_sum_indices = munkres_object.compute(t_matrix_large)

# Save the values of the maximum sum cells.
max_sum_indices_values = []

# Determine maximum sum.
for row, column in max_sum_indices:
    max_sum += t_matrix[row][column]
    max_sum_indices_values.append(t_matrix[row][column])

# Write the output file.
with open(output_file_path, 'w') as output_file:
    output_file.write(explanation)
    output_file.write("\nThe number of data points, n: " + str(n) + "\n")
    output_file.write("\nT(Ck, Lm) as a k by m matrix, where each row is a cluster, and each column is a class label:\n")
    output_file.write(str(t_matrix) + "\n")
    output_file.write("\nT(Ck, Lm) cells and values that make up the maximum sum:\n")
    for i in range(0, m):
        output_file.write(str(max_sum_indices[i]) + " = " + str(max_sum_indices_values[i]) + "\n")

    output_file.write("-------------\n\t" + str(max_sum) + "\n")
    output_file.write("\nAccuracy for clustering then equals this maximum sum divided by n." + "\n")
    output_file.write(str(max_sum) + "/" + str(n) + " = " + str(max_sum / n))

output_file.close()


import os
import json
import pandas as pd

def read_json_data(data_path : str, verbose=True):
    '''reads the contents of a json file, returns the content as a pandas dataframe to facilitate analysis
       json structure: {"data": [{"text": str, "label": str}]}'''

    # later on a check can be implemented to see if the parsed, tokenized and downsampled data is available.
    # but for the purpose of this task with small size, the performance is not affected by redoing it.

    json_path = __get_json_file_path(data_path)

    with open(json_path, encoding='utf8') as file:
        json_data = json.load(file)

    json_data = json_data['data']

    data_frame = pd.DataFrame.from_dict(json_data)

    if verbose:
        print('Parsed json file: ', len(data_frame), 'rows.')

    return data_frame

def store_datasets_as_csv(train : pd.DataFrame, dev : pd.DataFrame, test : pd.DataFrame, path_to_store : str):
    '''stores the provided pandas subsets as csv files'''

    if not os.path.isdir(path_to_store):
        os.makedirs(path_to_store)

    train.to_csv(os.path.join(path_to_store, 'train.csv'), header=False, index=False)
    dev.to_csv(os.path.join(path_to_store, 'dev.csv'), header=False, index=False)
    test.to_csv(os.path.join(path_to_store, 'test.csv'), header=False, index=False)

def __get_json_file_path(data_path):
    '''a helper method to get the name of the json file in the directory'''

    if not os.path.exists(data_path):
        raise ValueError("Data path doesn't exist.")

    json_files = [file for file in os.listdir(data_path) if file.endswith('.json')]

    if not len(json_files):
        raise ValueError('No json files in the provided path.')
    elif len(json_files) > 1:
        raise ValueError('Multiple json files exist. Please keep only one json file in this directory.')

    json_path = os.path.join(data_path, json_files[0])

    return json_path


import socket
import sys
import datetime

from threading import Thread
from time import sleep

from database.database import Database

class Server:
    def __init__(self):
        print("Data Server Created")
        self.received_data = []

    def encode(self, data):
        return data.encode("utf-8")

    def decode(self, data):
        return data.decode("utf-8")

    def request_measurements(self):
        while 1:
            now = datetime.datetime.now()
            file_name = "./storage/{}{}{}".format(now.day, now.month, now.year)
            db = Database(file_name)

            response = "ERROR: Response Not Defined"
            try:
                host = "192.168.0.13"  # location of PIS
                port = 8501  # PIS port
                request = "hi"  # irrelevant, PIS always returns everything
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.connect((host, port))
                sock.send(self.encode(request))
                response = sock.recv(1024)
                response = self.decode(response)
                db.insert(response)  # insert into DB
                # response is the readings, append them onto end of list
                self.received_data.append(response)
                sock.close()
            except OSError as os_error:
                print(os_error)
                response = "ERROR: Could Not Connect to Data Server at {}:{}".format(
                    host, port)

            sleep(1)

    def start_server(self):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if(len(sys.argv) > 1):
            port = sys.argv[1]
        else:
            port = 15000

        s.bind(('0.0.0.0', int(port)))
        print("Server Started on Port", port)
        s.listen(5)
        while 1:
            conn, addr = s.accept()
            conn_thread = Thread(target=self.handle_conn, args=(conn,))
            conn_thread.start()

    def handle_conn(self, conn):
        data = conn.recv(1024)
        data = self.decode(data)

        response = data
        args = data.split(" ")
        method_type = args[0]
        if(method_type == "GET_DATA"):
            property_ = args[1]  # voltage, current power etc
            if(property_ == "A"):  # get ALL
                quantity = int(args[2])  # how many
                total = len(self.received_data)
                if(quantity <= total):
                    response = ""
                    start = total - quantity
                    for i in range(start, total):
                        response = response + self.received_data[i]
                else:
                    response = "ERROR: Quantity cannot exceed total"
        response = self.encode(response)
        conn.send(response)


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from measure_point import MeasurePoint
from measure_diameter import MeasureDiameter

piece = 1

pickPoints = False
interpolateMissingValues = False
meanGraph = False
multipleMeasures = False

data = pd.read_csv("../path/to/file/" + str(piece) + ".csv", header=None, names=["x", "y", "z"])

measures = np.zeros((3, 10)) #number of measurements, number of measures to be taken for each measurement
means = []
std = []
maxMinDiff = []

pickedPoints = []

def checkPoint(point, startingPoint, measurePoints):
        if not startingPoint.defined:
            if startingPoint.isPoint(point[0], point[1]):
                startingPoint.setRealPoint(point[0], point[1])
        for measurePoint in measurePoints:
            if not measurePoint.defined:
                if measurePoint.isPoint(point[0], point[1]):
                    measurePoint.setRealPoint(point[0], point[1])

def checkDiameter(diameter, startingDiameter, measureDiameters):
    if not startingDiameter.defined:
        if startingDiameter.isDiameter(diameter[0], diameter[1], diameter[2]):
            startingDiameter.setRealDiameter(diameter[2])
    for measureDiameter in measureDiameters:
        if not measureDiameter.defined:
            if measureDiameter.isDiameter(diameter[0], diameter[1], diameter[2]):
                measureDiameter.setRealDiameter(diameter[2])

def onClick(event):
    if event.button == 1 and event.dblclick:
        point = (round(event.xdata, 4), round(event.ydata, 4))
        pickedPoints.append(point)
        print(f"Point {len(pickedPoints)} picked: {point}")

def scan(y):
    run = data[data['y'] == (y / 1000)]

    runZ = run['z'].copy()
    runZ[runZ == '***'] = np.nan
    xStart = np.array(run['x'], dtype=np.float64)
    zStart = np.array(runZ, dtype=np.float64) / 1000
    runMat = np.matrix([xStart, zStart])

    if not pickPoints:
        try:
            points = np.genfromtxt("points_" + str(piece) + ".csv", delimiter=',')
        except FileNotFoundError:
            print("Pick points before run")
            return

        measurePoints = []
        measureDiameters = []
        for i in range(len(points)):
            point = points[i]
            if i == 0:
                startingPoint = MeasurePoint(point[0], point[1])
            elif i >= 1 and i <= 7:
                measurePoints.append(MeasurePoint(point[0], point[1]))
            elif i >= 8 and i <= 17:
                if i % 2 == 0:
                    lastPointX = point[0]
                else:
                    if i == 9:
                        startingDiameter = MeasureDiameter(lastPointX, point[0], point[1])
                    else:
                        measureDiameters.append(MeasureDiameter(lastPointX, point[0], point[1]))
            elif i == 18:
                lastDiam1 = point[0]
            elif i == 19:
                lastDiam2 = point[0]

        refIndices = np.where(((lastDiam1 <= xStart) & (xStart <= lastDiam2)) | ((startingDiameter.xStartRef <= xStart) & (xStart <= startingDiameter.xEndRef)))[0]
        xHorizontalRef = xStart[refIndices]
        zHorizontalRef = zStart[refIndices]
        nonNanIndices = np.logical_not(np.isnan(zHorizontalRef))
        xHorizontalRef = xHorizontalRef[nonNanIndices]
        zHorizontalRef = zHorizontalRef[nonNanIndices]
        m, q = np.polyfit(xHorizontalRef, zHorizontalRef, 1)

        a = -np.arctan(m)
        rotMat = np.matrix([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
        rotatedRun = np.matmul(rotMat, runMat)

        run = pd.DataFrame(rotatedRun.transpose(), columns=['x', 'z'])
    else:
        run = pd.DataFrame(runMat.transpose(), columns=['x', 'z'])

    if interpolateMissingValues:
        dataGraph = run.copy()
        dataGraph = dataGraph.drop(dataGraph[np.isnan(dataGraph['z'])].index)
        xGraph = dataGraph['x']
        zGraph = np.array(dataGraph['z'], dtype=np.float64)
    else:
        xGraph = run['x']
        zGraphArray = run['z'].copy()
        zGraph = np.array(zGraphArray, dtype=np.float64)

    run = run.drop(run[np.isnan(run['z'])].index)

    runShift = run.shift(1)
    runShift.drop(runShift.index[:1], inplace=True)
    zShift = np.array(runShift['z'], dtype=np.float64)

    run = run.drop(run.index[:1])
    x = np.array(run['x'], dtype=np.float64)
    z = np.array(run['z'], dtype=np.float64)

    if not pickPoints:
        splitIndex = np.where((z < zShift - 0.1) | (z > zShift + 0.1))[0]
        splitIndex = np.insert(splitIndex, 0, 0)
        splitIndex = np.append(splitIndex, len(z))

        splitIndexDiameters = np.where((z < zShift - 0.02) | (z > zShift + 0.02))[0]
        if splitIndexDiameters[0] != 0:
            splitIndexDiameters = np.insert(splitIndexDiameters, 0, 0)
        if splitIndexDiameters[-1] != len(z):
            splitIndexDiameters = np.append(splitIndexDiameters, len(z))

        for i in range(len(splitIndex) - 1):
            xGroup = x[splitIndex[i]:splitIndex[i + 1]]
            zGroup = z[splitIndex[i]:splitIndex[i + 1]]
            checkPoint((xGroup[0], zGroup[0]), startingPoint, measurePoints)
            checkPoint((xGroup[-1], zGroup[-1]), startingPoint, measurePoints)

        zMeans = np.zeros(len(x))
        for i in range(len(splitIndexDiameters) - 1):
            zGroup = z[splitIndexDiameters[i]:splitIndexDiameters[i + 1]]
            mean = np.mean(zGroup)
            zMeans[splitIndexDiameters[i]:splitIndexDiameters[i + 1]] = mean
            xGroup = x[splitIndexDiameters[i]:splitIndexDiameters[i + 1]]
            checkDiameter([xGroup[0], xGroup[-1], mean], startingDiameter, measureDiameters)

        if startingPoint.defined:
            for i, measurePoint in enumerate(measurePoints):
                if measurePoint.defined:
                    measure = startingPoint.x - measurePoint.x
                    measures[y][i] = measure
                else:
                    print(f"Measure point {i + 1} not found for y={y / 1000}")

        if startingDiameter.defined:
            for i, measureDiameter in enumerate(measureDiameters):
                if measureDiameter.defined:
                    diameter = abs(measureDiameter.zMean - startingDiameter.zMean)
                    measures[y][len(measurePoints) + i] = diameter
                else:
                    print(f"Measure diameter {i + 1} not found for y={y / 1000}")

        if not multipleMeasures:
            print(measures[y])

            if meanGraph:
                xGraph = x
                zGraph = zMeans

            plt.plot(xGraph, zGraph)
            plt.xlim(np.nanmax(xGraph) + 3, np.nanmin(xGraph) - 3)
            plt.ylim(np.nanmin(zGraph) - 1, np.nanmax(zGraph) + 1)
            plt.show()

    else:
        fig, ax = plt.subplots()
        ax.plot(xGraph, zGraph)
        ax.set_xlim(np.nanmax(xGraph) + 3, np.nanmin(xGraph) - 3)
        ax.set_ylim(np.nanmin(zGraph) - 1, np.nanmax(zGraph) + 1)
        fig.canvas.mpl_connect('button_press_event', onClick)
        plt.show()
        if len(pickedPoints) == 20:
            np.savetxt("points_" + str(piece) + ".csv", pickedPoints, delimiter=',')
            print("Points file saved")

def main():
    if multipleMeasures:
        if pickPoints:
            print("PickPoints and MultipleMeasures cannot be used together")
            return

        for y in range(measures.shape[0]):
            scan(y)

        #np.savetxt("measures.csv", measures, delimiter=';')

        for dim in range(measures.shape[1]):
            allMeasures = measures[:, dim]
            allMeasures = allMeasures[allMeasures != 0]
            means.append(np.mean(allMeasures))
            std.append(np.std(allMeasures, ddof=1))
            maxMinDiff.append(np.max(allMeasures) - np.min(allMeasures))

        print("Means", means)
        print("Std", std)
        print("MaxMinDiff", maxMinDiff)
    else:
        scan(0) #editable

if __name__ == "__main__":
    main()


import pandas as pd
import geopandas as gpd
import numpy as np
import matplotlib.pyplot as plt
import sys
import folium
import branca
from features import *
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import cm
from folium.plugins import MeasureControl

class mapfun(object):

	def __init__(self, x:float, y:float, zoom:int=10):
		self.map = folium.Map(location=[y,x], zoom_start=zoom)

	def add_tiles(self, tile_type:str='World_Imagery', services:str=''):
		'''
		add different tiles to the interactive map
		inputs
			folmap = intitialized interactive map
			Defaults:
				- tile type = esri tile layer name (World_Imagery, USA_Topo_Maps, World_Shaded_Relief,
													World_Street_Map)
				- services  = access service sub dir or choose other ones (/Canvas,/Elevation,/Ocean,
																			/Polar, etc)
				*note*  : if you change services then the tile needs to be changed according to the service
							you want.
		'''
		EsriImagery = "https://server.arcgisonline.com/ArcGIS/rest/services"\
					f"{services}/{tile_type}/MapServer/"\
					"tile/{z}/{y}/{x}"
		EsriAttribution = "Tiles © Esri — Source: Esri, i-cubed,"\
					" USDA, USGS, AEX, GeoEye, Getmapping, Aerogrid, IGN, IGP"\
					", UPR-EGP, and the GIS User Community"
		folium.TileLayer(tiles=EsriImagery, attr=EsriAttribution, name=tile_type).add_to(self.map)
		return self

	def add_cm_points(self, gdf0:gpd.GeoDataFrame, info:list, name:list, cmap:list, add_table:bool=True):
		'''
		 adds points with colormap to folium map
		 inputs:
		 	gdf0:				geopandas.GeoDataFrame
			cmap:				list of colors
			add_table :         choose if you would like to add a table to points
			info:  				adds descriptions in table format
								to polygons
			name:				list of the names of the polygons

		'''

		if add_table:
			for idx in  gdf0.index:
				iframe=branca.element.IFrame(html=info[idx], width=box_width(info[idx]),
											height=box_height(info[idx]))
				popup = folium.Popup(iframe)
				self.map.add_child(folium.CircleMarker(location=[gdf0.loc[idx]['geometry'].y,
											 gdf0.loc[idx]['geometry'].x],
											 popup= popup,radius=4, fill=True, weight=4,
											  color=mpl.colors.to_hex(cmap),
											  fill_color=mpl.colors.to_hex(cmap), fill_opacity=1))
		else:
			for idx in  gdf0.index:
				self.map.add_child(folium.CircleMarker(location=[gdf0.loc[idx]['geometry'].y,
											 gdf0.loc[idx]['geometry'].x],
											 radius=4, fill=True, weight=4,
											  color=mpl.colors.to_hex(cmap),
											  fill_color=mpl.colors.to_hex(cmap), fill_opacity=1))
		return self

	def add_point(self, gdf0:gpd.GeoDataFrame, info:list, name:list, color:str='red',\
			fillcolor:str='red',data:list=None,add_table:bool=True):
		'''
		 adds points to folium map
		 inputs:
		 	gdf0:				geopandas.GeoDataFrame

			add_table :         choose if you would like to add a table to points
			info:  				adds descriptions in table format
								to points
			name:				list of the names of the points

		'''
		sID=name

		if type(color) == list:
			color = colormap(data)
			fillcolor = colormap(data)
		if add_table:
			for idx in  gdf0.index:
				iframe=branca.element.IFrame(html=info[idx], width=box_width(info[idx]),
											height=box_height(info[idx]))
				popup = folium.Popup(iframe)
				self.map.add_child(folium.CircleMarker(location=[gdf0.loc[idx]['geometry'].y,
											 gdf0.loc[idx]['geometry'].x],
											 popup= popup,radius=4, fill=True, weight=4,
											  color=color,fill_color=fillcolor, fill_opacity=1))
		else:
			for idx in  gdf0.index:
				self.map.add_child(folium.CircleMarker(location=[gdf0.loc[idx]['geometry'].y,
											 gdf0.loc[idx]['geometry'].x],
											 radius=4, fill=True, weight=4,
											  color=color,fill_color=fillcolor, fill_opacity=1))
		return self

	def add_lines(self, gdf0:gpd.GeoDataFrame, info:list, name:list,\
				 color:str='blue',fillcolor:str='blue',add_table:bool=True):
		'''
		 adds lines to folium map
		 inputs:
		 	gdf0:				geopandas.GeoDataFrame

			add_table :         choose if you would like to add a table to points
			info:  				adds descriptions in table format
								to lines
			name:				list of the names of the lines

		'''
		sID=name
		if add_table:
			for idx in  gdf0.index:
				geometry =  gdf0.loc[idx].geometry
				if geometry.type == 'MultiLineString':

					iframe=branca.element.IFrame(html=info[idx], width=box_width(info[idx]),
						height=box_height(info[idx]))
					popup = folium.Popup(iframe)
					for line in gdf0.geometry[idx]:
						x = line.coords.xy[0]
						y = line.coords.xy[1]
						line_points = zip(y,x)
						style_function = lambda x :{'fillColor': fillcolor,'color': color,'opacity':0.6,'fillOpacity': 0.4}
						geojson = folium.GeoJson(line,name=name[idx],style_function=style_function).add_to(self.map)
						geojson.add_child(popup)
				else:
					iframe=branca.element.IFrame(html=info[idx], width=box_width(info[idx]),
												height=box_height(info[idx]))
					popup = folium.Popup(iframe)

					#assert geometry.type  != 'MultiLineString', "MultiLineString in dataset, convert to line in GIS before continuing"

					x = geometry.coords.xy[0]
					y = geometry.coords.xy[1]
					style_function = lambda x :{'fillColor': fillcolor,'color': color,'opacity':0.6,'fillOpacity': 0.4}
					line_points = zip(y,x)
					geojson = folium.GeoJson(gdf0.geometry[idx],name=name[idx],style_function=style_function).add_to(self.map)
					geojson.add_child(popup)
		else:
			for idx in  gdf0.index:

				geometry =  gdf0.loc[idx].geometry
				if geometry.type == 'MultiLineString':

					for line in gdf0.geometry[idx]:
						x = line.coords.xy[0]
						y = line.coords.xy[1]
						line_points = zip(y,x)
						style_function = lambda x :{'fillColor': fillcolor,'color': color,'opacity':0.6,'fillOpacity': 0.4}
						geojson = folium.GeoJson(line,name=name[idx],style_function=style_function).add_to(self.map)

				else:
					x = geometry.coords.xy[0]
					y = geometry.coords.xy[1]
					style_function = lambda x :{'fillColor': fillcolor,'color': color,'opacity':0.6,'fillOpacity': 0.4}
					line_points = zip(y,x)
					geojson = folium.GeoJson(gdf0.geometry[idx],name=name[idx],style_function=style_function).add_to(self.map)

		return self

	def add_polygon(self, gdf0:gpd.GeoDataFrame, info:list, name:list,\
				 color:str='black',fillcolor:str='black',add_table:bool=True):
		'''
		 adds polygons to folium map
		 inputs:
		 	gdf0:				geopandas.GeoDataFrame

			add_table :         choose if you would like to add a table to points
			info:  				adds descriptions in table format
								to polygons
			name:				list of the names of the polygons

		'''
		if add_table:

			for idx in  gdf0.index:
				iframe=branca.element.IFrame(html=info[idx], width=box_width(info[idx]),
											height=box_height(info[idx]))
				popup = folium.Popup(iframe)
				style_function = lambda x :{'fillColor': fillcolor,'color': color,'opacity':0.6,'fillOpacity': 0.4}

				geojson = folium.GeoJson(gdf0.geometry[idx],name=name[idx],style_function=style_function).add_to(self.map)
				geojson.add_child(popup)
		else:
			for idx in  gdf0.index:

				style_function = lambda x :{'fillColor': fillcolor,'color': color,'opacity':0.6,'fillOpacity': 0.4}

				geojson = folium.GeoJson(gdf0.geometry[idx],name=name[idx],style_function=style_function).add_to(self.map)
		return self

	def add_extra(self):
		folium.LayerControl().add_to(self.map)
		self.map.add_child(MeasureControl())
		return self

#---------- extra attribute functions that support the functions above---------------

def colormap(data:list, cmap_type:str='jet', vmin=0):
	'''
	 calculates colors based on the data and cmap
	'''
	cmap = cm.get_cmap(cmap_type)
	normalize = mpl.colors.Normalize(vmin=vmin, vmax=max(data))
	colors = [cmap(normalize(value)) for value in data]

	return colors

def colors():
	'''
	List of colors for branca colorbar
	needs to be updated 7/12
	'''
	colors = ['#8000FF','#5500FF','#4000FF','#1500FF','#0000FF','#002AFF','#0055FF',
	'#0080FF','#00AAFF','#00D4FF','#00FFFF','#00FFAA','#00FF80','#00FF2A','#2AFF00',
	'#80FF00','#AAFF00','#D4FF00','#FFFF00','#FFD400','#FFAA00','#FF8000','#FF6A00',
	'#FF4000','#FF2A00','#FF1500','#FF0000']
	return colors

def add_cbar(data:list,folmap:folium.folium.Map,caption:str='HWM(ft) above NAVD88', vmin=0):
	'''
	adds a colorbar to a folium map
	'''
	colormap = branca.colormap.StepColormap(colors=colors(),vmin=vmin,vmax=np.max(round(data,1)))
	colormap.caption = caption
	colormap.add_to(folmap)

	return folmap

def box_width(descriptions:str):
	'''
	automatically sets the width of the html box based on the inputs given
	'''
	width = 175
	num_columns = find_all(descriptions,'COL width=')
	for i in num_columns:
		width += int(descriptions[i:(i+20)].split('"')[1].split('%')[0])
	return width

def box_height(descriptions:str):
	'''
	automatically adjusts the height of the html box
	'''
	height = 100
	num_rows = find_all(descriptions,'<TR>')
	height = height+(len(num_rows)*20)
	return height

def quick_transform(vector, incrs):
	'''
	sets the crs projection for the vector
	'''
	vector.crs = incrs
	return vector.to_crs({'init': 'epsg:4326'})

def get_centroid(vector):
	'''
	finds the center points in x,y from a geopandas.GeoDataFrame
	'''
	geom = vector.geometry
	xpoints = [float(v.x) for v in geom]
	ypoints = [float(v.y) for v in geom]
	return np.mean(xpoints), np.mean(ypoints)


'''
Early Access Training Code
'''
import warnings
warnings.simplefilter("ignore", UserWarning)

import os
import argparse
from tqdm import tqdm
import matplotlib.pyplot as plt

from einops import rearrange, reduce, repeat
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
import torch.multiprocessing as mp

from pytorch3d.renderer import (
    MonteCarloRaysampler,
    GridRaysampler,
    PerspectiveCameras
)

from utils.camera_utils import RelativeCameraLoader
from utils.common_utils import huber, sample_images_at_mc_locs
from utils.render_utils import CustomImplicitRenderer
from utils.render_utils import LightFieldRaymarcher
from utils.load_model import load_vae
from utils.co3d_dataloader import CO3Dv2Wrapper

from external.imagen_pytorch import Unet
from sparsefusion.eft import EpipolarFeatureTransformer
from sparsefusion.vldm import DDPM

def save_visualization(
    eft, diffusion, vae, scale_factor, camera,
    target_image,
    loss_history_color,
    renderer_grid,
    dir='output/test',
    mod='',
    input_rgb=None,
    input_cameras=None,
    z=None,
    args=None
    ):
    """
    Save visualization
    """
    if not os.path.exists(dir):
        os.makedirs(dir)

    with torch.no_grad():

        if args.train_eft:
            epipolar_features , _, _ = renderer_grid(
                cameras=camera,
                volumetric_function=eft.module.batched_forward
            )
        else:
            epipolar_features , _, _ = renderer_grid(
                cameras=camera,
                volumetric_function=eft.batched_forward
            )

        rendered_image, epipolar_latents = (
            epipolar_features.split([3, 256], dim=-1)
        )

        epipolar_latents = rearrange(epipolar_latents, 'b h w f -> b f h w')
        rand = torch.randn_like(z)
        diffusion_z = diffusion.module.sample(cond_images=epipolar_latents, batch_size=1)
        diffusion_z = 1.0 / args.z_scale_factor * diffusion_z
        diffusion_image = unnormalize(vae.decode(diffusion_z)).clip(0.0, 1.0)
        diffusion_image = rearrange(diffusion_image, 'b c h w -> b h w c')

    num_input = len(input_rgb)
    # Generate plots.
    fig, ax = plt.subplots(1, num_input+3, figsize=(2*(num_input+3), 3))
    ax = ax.ravel()
    clamp_and_detach = lambda x: x.clamp(0.0, 1.0).cpu().detach().numpy()
    for ai in range(num_input):
        ax[ai].imshow(clamp_and_detach(input_rgb[ai]))
        ax[ai].grid("off")
        ax[ai].axis("off")
        ax[ai].set_title('input')

    ax[-3].imshow(clamp_and_detach(target_image[0]))
    ax[-3].grid("off")
    ax[-3].axis("off")
    ax[-3].set_title('target')
    ax[-2].imshow(clamp_and_detach(rendered_image[0]))
    ax[-2].grid("off")
    ax[-2].axis("off")
    ax[-2].set_title('pred')
    ax[-1].imshow(clamp_and_detach(diffusion_image[0]))
    ax[-1].grid("off")
    ax[-1].axis("off")
    ax[-1].set_title('pred')

    plt.savefig(f'{dir}/{mod}.jpg', bbox_inches='tight')
    plt.cla()
    plt.close(fig)

    plt.plot(list(range(len(loss_history_color))), loss_history_color, linewidth=1)
    plt.savefig(f'{dir}/_loss.jpg')
    plt.cla()
    plt.close(fig)

    plt.plot(list(range(len(loss_history_color[-100:]))), loss_history_color[-100:], linewidth=1)
    plt.savefig(f'{dir}/_loss_recent.jpg')
    plt.cla()
    plt.close(fig)

def vis_helper(step, args, eft, diffusion, vae, renderer_grid, target_cameras, target_rgb, relative_cam, loss_history, z):
    # Visualize the full renders every 100 iterations.
    '''
    target_cameras: ()
    target_rgb: (B, C, H, W) or (B, H, W, C) but we want (B, H, W, C)
    '''
    if target_rgb.shape[-1] != 3 and len(target_rgb.shape) == 4:
        target_rgb = target_rgb.permute(0,2,3,1) # (B, C, H, W)

    #! SAMPLE BATCH CAMERAS
    rand_batch = torch.randperm(len(target_cameras))
    batch_idx = rand_batch[:1]
    if args.num_input_range is not None:
        context_size = torch.randint(args.num_input_range[0], args.num_input_range[1],(1,))
    else:
        context_size = args.num_input
    render_camera, render_rgb, batch_mask, input_cameras, input_rgb, input_masks = relative_cam(target_cameras, target_rgb, context_size=context_size, query_idx=batch_idx)

    with torch.no_grad():

        if args.train_eft:
            eft.module.encode(input_cameras, input_rgb)
        else:
            eft.encode(input_cameras, input_rgb)

        save_visualization(
            eft,
            diffusion,
            vae,
            args.scale_factor,
            render_camera,
            render_rgb,
            loss_history,
            renderer_grid,
            dir=f'{args.exp_dir}/log/',
            mod=f'train_{step:06d}',
            input_rgb=input_rgb,
            input_cameras=input_cameras,
            z=z,
            args=args
        )

def normalize(x):
    return torch.clip(x*2 - 1.0, -1.0, 1.0)

def unnormalize(x):
    return torch.clip((x + 1.0) / 2.0, 0.0, 1.0)

def init_renderer(gpu, img_h, img_w, min=0.1, max=4.0, bbox=None, n_pts_per_ray=20, n_rays=750, scale_factor=None):
    '''
    Return renderers

    renderer_grid - full resolution
    renderer_feat - scaled down feature grid resolution
    renderer_mc - random sampling
    '''
    img_h, img_w = img_h, img_w
    volume_extent_world = max
    half_pix_width = 1.0 / img_w
    half_pix_height = 1.0 / img_h

    raysampler_grid = GridRaysampler(
        min_x=1.0 - half_pix_width,
        max_x=-1.0 + half_pix_width,
        min_y=1.0 - half_pix_height,
        max_y=-1.0 + half_pix_height,
        image_height=img_h,
        image_width=img_w,
        n_pts_per_ray=n_pts_per_ray,
        min_depth=min,
        max_depth=volume_extent_world,
    )
    if scale_factor is not None:
        raysampler_features = GridRaysampler(
            min_x=1.0 - half_pix_width,
            max_x=-1.0 + half_pix_width,
            min_y=1.0 - half_pix_height,
            max_y=-1.0 + half_pix_height,
            image_height=int(img_h//scale_factor),
            image_width=int(img_w//scale_factor),
            n_pts_per_ray=n_pts_per_ray,
            min_depth=min,
            max_depth=volume_extent_world,
        )
    if bbox is None:
        raysampler_mc = MonteCarloRaysampler(
            min_x = -1.0,
            max_x = 1.0,
            min_y = -1.0,
            max_y = 1.0,
            n_rays_per_image=n_rays,
            n_pts_per_ray=n_pts_per_ray,
            min_depth=min,
            max_depth=volume_extent_world,
        )
    elif bbox is not None:
        raysampler_mc = MonteCarloRaysampler(
            min_x = -bbox[0,1],
            max_x = -bbox[0,3],
            min_y = -bbox[0,0],
            max_y = -bbox[0,2],
            n_rays_per_image=n_rays,
            n_pts_per_ray=n_pts_per_ray,
            min_depth=min,
            max_depth=volume_extent_world,
        )

    raymarcher = LightFieldRaymarcher()

    renderer_grid = CustomImplicitRenderer(
        raysampler=raysampler_grid, raymarcher=raymarcher, reg=True
    )
    renderer_mc = CustomImplicitRenderer(
        raysampler=raysampler_mc, raymarcher=raymarcher, reg=True
    )

    renderer_grid = renderer_grid.cuda(gpu)
    renderer_mc = renderer_mc.cuda(gpu)

    if scale_factor is None:
        return renderer_grid, renderer_mc
    else:
        renderer_feat = CustomImplicitRenderer(
            raysampler=raysampler_features, raymarcher=raymarcher, reg=True
        )
        renderer_feat = renderer_feat.cuda(gpu)
        return renderer_grid, renderer_feat, renderer_mc

def load_dataset(args, image_size=256):
    '''
    Return dataset
    '''
    if args.dataset_name == 'co3d':
        train_dataset = CO3Dv2Wrapper(root=args.root, category=args.category, sample_batch_size=20, image_size=image_size)
    return train_dataset

def train(gpu, args):

    #! INIT DISTRIBUTED PROCESS
    rank = args.nr * args.gpus + gpu
    print('spawning gpu rank', rank, 'out of', args.gpus, 'using', args.backend)
    dist.init_process_group(args.backend, rank=rank, world_size=args.world_size)
    torch.cuda.set_device(gpu)
    os.makedirs(args.exp_dir, exist_ok=True)
    os.makedirs(args.exp_dir + '/log/', exist_ok=True)

    #! LOAD MODEL
    eft, vae, diffusion = load_model(gpu, args)

    #! OPTIMIZER
    lr = 5e-5
    optimizer = torch.optim.Adam(diffusion.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.5)
    diffusion = nn.parallel.DistributedDataParallel(diffusion, device_ids=[gpu], find_unused_parameters=True)

    if args.train_eft:
        print('training eft')
        optimizer_nerf = torch.optim.Adam(eft.parameters(), lr=lr)
        scheduler_nerf = torch.optim.lr_scheduler.StepLR(optimizer_nerf, step_size=50000, gamma=0.5)
        eft = nn.parallel.DistributedDataParallel(eft, device_ids=[gpu])

    #! LOAD DATALOADER
    batch_size = 1
    print('preparing to load ds on gpu', gpu)
    train_dataset = load_dataset(args, image_size=args.image_size)
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
                                                                    num_replicas=args.world_size,
                                                                    rank=rank)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                            batch_size=batch_size,
                                            shuffle=False,
                                            num_workers=4,
                                            pin_memory=True,
                                            sampler=train_sampler)
    print('training data len', len(train_dataset))

    #! INIT RELATIVE CAMERA LOADER
    relative_cam = RelativeCameraLoader(relative=True, center_at_origin=True)

    #! INITIALIZE TRAINING LOOP
    #! TRAIN PARAMTERS
    n_repeat = args.repeat
    max_epoch = 1000
    vis_itr=args.vis_itr
    save_itr=args.save_itr
    render_batch_size = 1

    max_step = n_repeat*max_epoch*len(train_loader)
    if gpu == 0:
        pbar = tqdm(total=max_step)
        step = 0
        pbar.n = step
        pbar.last_print_n = step
        pbar.refresh()

    loss_history = []
    running_loss = 0
    for ep in range(max_epoch):
        for data in train_loader:

            for re in range(n_repeat):
                #! FIT CURRENT BATCH
                optimizer.zero_grad()
                if args.train_eft:
                    optimizer_nerf.zero_grad()
                torch.autograd.set_detect_anomaly(True)

                #! LOAD DATA
                target_cameras = PerspectiveCameras(R=data['R'][0],T=data['T'][0],focal_length=data['f'][0],principal_point=data['c'][0],image_size=data['image_size'][0]).cuda(gpu)
                target_rgb = data['images'][0].cuda(gpu)
                bbox = data['bbox'][0].cuda(gpu)
                valid_region = data['valid_region'][0].cuda(gpu)

                #! SAMPLE BATCH CAMERAS
                rand_batch = torch.randperm(len(target_cameras))
                batch_idx = rand_batch[:render_batch_size]
                if args.num_input_range is not None:
                    context_size = torch.randint(args.num_input_range[0], args.num_input_range[1],(1,))
                else:
                    context_size = args.num_input
                batch_cameras, batch_rgb, batch_mask, input_cameras, input_rgb, input_masks, context_idx = relative_cam(target_cameras, target_rgb, context_size=context_size, query_idx=batch_idx, return_context=True)
                input_bbox = bbox[context_idx]
                batch_valid_region = valid_region[batch_idx]

                #! INITIALIZE RENDERERS
                if args.dataset_name == 'co3d':
                    cam_dist_mean = torch.mean(torch.linalg.norm(target_cameras.get_camera_center(), axis=1))
                    min_depth = cam_dist_mean - 5.0
                    volume_extent_world = cam_dist_mean + 5.0
                    if bbox is None:
                        renderer_grid, renderer_feat, renderer_mc = init_renderer(gpu, train_dataset.img_h, train_dataset.img_w, min=min_depth, max=volume_extent_world, scale_factor=args.scale_factor)
                    elif bbox is not None:
                        renderer_grid, renderer_feat, renderer_mc = init_renderer(gpu, train_dataset.img_h, train_dataset.img_w, min=min_depth, max=volume_extent_world, bbox=bbox[batch_idx], scale_factor=args.scale_factor)

                #! EFT
                if args.train_eft:
                    epipolar_features, sampled_rays, reg_term = renderer_feat(
                        cameras=batch_cameras,
                        volumetric_function=eft.module.batched_forward,
                        n_batches=1,
                        input_cameras=input_cameras,
                        input_rgb=input_rgb
                    )
                else:
                    with torch.no_grad():
                        epipolar_features, sampled_rays, reg_term = renderer_feat(
                            cameras=batch_cameras,
                            volumetric_function=eft.batched_forward,
                            n_batches=1,
                            input_cameras=input_cameras,
                            input_rgb=input_rgb
                        )

                rendered_images, epipolar_latents = (
                    epipolar_features.split([3, 256], dim=-1)
                )

                #! RESHAPE
                diffusion_batch_size = args.diffusion_batch_size
                if batch_rgb.shape[1] != 3:
                    batch_rgb = rearrange(batch_rgb, 'b h w c -> b c h w')
                if args.dataset_name == 'srn':
                    batch_rgb = F.interpolate(batch_rgb, scale_factor=2.0)
                with torch.no_grad():
                    images_z = vae.encode(normalize(batch_rgb)).mode() * args.z_scale_factor

                epipolar_latents = rearrange(epipolar_latents, 'b h w f -> b f h w')

                diffusion_input = images_z.expand(diffusion_batch_size, -1, -1, -1)
                diffusion_cond_image = epipolar_latents.expand(diffusion_batch_size, -1, -1, -1)

                #! MASK LOSS OUTSIDE OF VALID REGION
                batch_mask = F.interpolate(batch_valid_region, scale_factor=0.125, mode='bilinear')
                batch_mask = batch_mask.expand(diffusion_batch_size, images_z.shape[1], -1, -1)
                batch_mask[batch_mask > 0.6] = 1.0
                batch_mask[batch_mask <= 0.6] = 0.0

                color_loss = 0
                if args.train_eft:
                    #! COMPUTE COLOR ERROR
                    colors_at_rays = sample_images_at_mc_locs(
                        batch_rgb,
                        sampled_rays.xys
                    )

                    color_loss = huber(
                        rendered_images,
                        colors_at_rays,
                    )

                    color_loss = color_loss * rearrange(batch_mask[:1, :1, ...], 'b c h w -> b h w c')
                    color_loss = color_loss.abs().mean()

                #! DIFFUSION ERROR
                d_loss = diffusion(diffusion_input, cond_images=diffusion_cond_image, loss_mask=batch_mask)

                #! LOSS TERM
                loss = d_loss + color_loss

                #! BACKWARD
                #! CATCH SOME ERRORS
                try:
                    loss.backward()
                    skip_save = False
                except RuntimeError:
                    print('RUNTIME ERROR')
                    print(data['idx'])
                    return

                #! OPTIM STEP
                optimizer.step()
                scheduler.step()
                if args.train_eft:
                    optimizer_nerf.step()
                    scheduler_nerf.step()

                if gpu == 0:
                    step += 1
                    pbar.update(1)
                    pbar.set_description(f'epoch={ep}/{max_epoch} lr={float(scheduler.get_last_lr()[0]):1.2e} loss={loss.item():02f}')
                    running_loss += loss.item()

                    #! UPDATE LR AND LOG LOSS
                    if step % 50 == 0:
                        loss_history.append(float(running_loss)/100)
                        running_loss = 0

                    #! VISUALIZE
                    if step % vis_itr == 0:
                        print('visualizing', args.exp_name, args.category)
                        vis_helper(step, args, eft, diffusion, vae, renderer_feat, target_cameras, target_rgb, relative_cam, loss_history, z=images_z)

                    #! SAVE MODEL
                    if step % save_itr == 0 and step!=0 and not skip_save:
                        save_model(args, step, f'latest', diffusion, optimizer, scheduler, eft)
                        if step % 50000 == 0:
                            save_model(args, step, f'{step:06d}', diffusion, optimizer, scheduler, eft)
                            # save_model(args, step, f'prev', diffusion, optimizer, scheduler, eft)
                        print('saving model at step', step)

def save_model(args, step, mod, model, optimizer, scheduler, eft=None):
    '''
    Saves Models
    '''
    torch.save({
        'step': step,
        'model_state_dict': model.module.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        }, f'{args.exp_dir}/ckpt_{mod}.pt')
    if args.train_eft:
        torch.save({
        'step': step,
        'model_state_dict': eft.module.state_dict(),
        }, f'{args.exp_dir}/ckpt_{mod}_eft.pt')

def load_model(gpu, args):
    '''
    Loads Models
    '''
    #! LOAD EFT (eft)
    eft = EpipolarFeatureTransformer(use_r=args.use_r, encoder=args.encoder, return_features=True, remove_unused_layers=False).cuda(gpu)

    if args.ckpt_path is not None:
        checkpoint = torch.load(args.ckpt_path, map_location='cpu')
        eft.load_state_dict(checkpoint['model_state_dict'])
    print('LOADING 1/3 loaded eft checkpoint from', args.ckpt_path)

    #! LOAD VAE
    vae = load_vae(args.ckpt_path_vae).cuda(gpu)
    print('LOADING 2/3 loaded sd vae from', args.ckpt_path_vae)

    #! LOAD UNet
    channels = 4
    feature_dim = 256
    unet1 = Unet(
        channels=channels,
        dim = 256,
        dim_mults = (1, 2, 4, 4),
        num_resnet_blocks = (2, 2, 2, 2),
        layer_attns = (False, False, False, True),
        layer_cross_attns = (False, False, False, False),
        cond_images_channels = feature_dim,
        attn_pool_text=False
    )
    total_params = sum(p.numel() for p in unet1.parameters())
    print(f"{unet1.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")

    #! LOAD DIFFUSION
    print('using diffusion objective', args.objective)
    diffusion = DDPM(
        channels=channels,
        unets = (unet1, ),
        conditional_encoder = None,
        conditional_embed_dim = None,
        image_sizes = (32, ),
        timesteps = 500,
        cond_drop_prob = 0.1,
        pred_objectives=args.objective,
        conditional=False,
        auto_normalize_img=False,
        clip_output=True,
        dynamic_thresholding=False,
        dynamic_thresholding_percentile=.68,
        clip_value=10,
    ).cuda(gpu)

    if args.ckpt_path_diffusion is not None:
        checkpoint = torch.load(args.ckpt_path_diffusion, map_location='cpu')
        diffusion.load_state_dict(checkpoint['model_state_dict'])
        print('LOADING 3/3 loaded diffusion from', args.ckpt_path_diffusion)
    else:
        print('LOADING 3/3 loaded diffusion from', 'scratch')

    return eft, vae, diffusion

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
    parser.add_argument('-g', '--gpus', default=1, type=int,
                        help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int,
                        help='ranking within the nodes')
    parser.add_argument('-p', '--port', default=1, type=int, metavar='N',
                        help='last digit of port (default: 1234[1])')
    parser.add_argument('-c', '--category', type=str, metavar='s', required=True,
                        help='category')
    parser.add_argument('-r', '--root', type=str, default='/grogu/datasets/co3d/', metavar='s',
                        help='location of test features')
    parser.add_argument('-d', '--dataset_name', type=str, default='co3d', metavar='s',
                        help='dataset name')
    parser.add_argument('-b', '--backend', type=str, default='nccl', metavar='s',
                        help='nccl')
    parser.add_argument('-a', '--vae', type=str, default='checkpoints/sd/sd-v1-3-vae.ckpt', metavar='S',
                        help='vae ckpt')
    args = parser.parse_args()
    args.world_size = args.gpus * args.nodes
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = f'1234{args.port}'
    print('using port', f'1234{args.port}')

    torch.manual_seed(1)

    if args.dataset_name == 'co3d':

        #! ####################################################################
        #! TODO:
        #! ####################################################################

        #! EXPERIMENT GROUP FOLDER
        args.exp_group = 'sparsefusion_exp'

        #! EXPERIMENT NAME
        args.exp_name = 'demo_train'

        #! EXPERIMENT MODIFIDER
        args.mod = '_joint'

        #! EXPERIMENT DIRECTORY
        args.exp_dir = f'output/{args.exp_group}/{args.exp_name}/{args.category}{args.mod}'

        #! STABLE DIFFUSION VAE PATH
        args.ckpt_path_vae = args.vae

        #! NUMBER OF INPUTS VIEWS
        args.num_input_range = (2,6) # [low, high)

        #! ####################################################################
        #! ABOVE: TODO
        #! ####################################################################

        #@ AUTOMATIC RESUME
        #! WARNING: CURRENTLY DOES NOT RESUME OPTIMIZER
        if os.path.exists(f'{args.exp_dir}/ckpt_latest_eft.pt'):
            print('***automatically resuming***')
            args.ckpt_path = f'{args.exp_dir}/ckpt_latest_eft.pt'
            args.ckpt_path_diffusion = f'{args.exp_dir}/ckpt_latest.pt'
        else:
            print('***training from scratch***')
            args.ckpt_path = None # or default pretrained
            args.ckpt_path_diffusion = None # or default pretrained

        #@ SET PARAMETERS
        args.timesteps = 500                # default diffusion steps (500)
        args.objective = 'noise'            # default objective (noise)
        args.scale_factor = 8               # default VAE scale factor
        args.image_size = 256               # default image resolution
        args.dataset_name = 'co3d'          # dataset name
        args.diffusion_batch_size = 12      # extend batch size (12)
        args.repeat = 1                     # repeat a scene (1)
        args.vis_itr = 100                  # default visualization freq (100)
        args.save_itr = 1000                # save iteration
        args.train_eft = True               # jointly train EFT (eft)
        if args.train_eft:
            print('joint eft training')

    # Modification below not recommended
    args.use_r = True
    args.encoder = 'resnet18'
    args.num_input = 4
    args.z_scale_factor = 0.18215

    mp.spawn(train, nprocs=args.gpus, args=(args,))

if __name__ == '__main__':
    os.environ["TORCH_DISTRIBUTED_DEBUG"]='INFO'
    main()


import django_filters
from django import forms
from django_filters import FilterSet
from .models import *

class PostFilter(FilterSet):
    class Meta:
        model = Post
        fields = {
            'title': ['icontains'],
            'post_text': ['icontains'],
            'date_create': ['gt', 'lt'],

        }
        filter_overrides = {
            models.DateField: {
                'filter_class': django_filters.DateFilter,
                'extra': lambda f: {
                    'widget': forms.DateInput,
                },
            },
        }


import cv2
import numpy as np
import os
import face_recognition
import json

from multiprocess import pool

import csv

def savecsv(toCSV) :

    csv_columns =toCSV[0].keys()
    csv_file="Datadet.csv"
    try:
        with open(csv_file, 'w') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
            writer.writeheader()
            for data in toCSV:
                writer.writerow(data)
    except IOError:
        print("I/O error")

def regr(ks , camera) :
    if not os.path.exists("data"):
        os.makedirs("data")   # this is the location of pictures of the students , used for backup

    if not os.path.exists("json"):
        os.makedirs("json")   # all json files will be stored here

    direc = "data/"
    cap = cv2.VideoCapture(camera)  #'http://192.168.137.163:4747/mjpegfeed'  it changes every time a connectipn is made

    while True :
        _, frame = cap.read()
     # frame size
        x1 = int(10)
        y1 = 10
        x2 = int(frame.shape[1])
        y2 = int(frame.shape[1])

        cv2.rectangle(frame, (x1-1, y1-1), (x2+1, y2+1), (255,0,0) ,1)

        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.imshow("Frame", frame)

        interrupt = cv2.waitKey(10)

        if  interrupt & 0xFF ==ord('1') :

            try :
                fi = face_recognition.face_encodings(frame)[0]
                cv2.imwrite(direc +ks+'.jpg', frame)
                cap.release()
                cv2.destroyAllWindows()
                return frame , fi

            except IndexError :
                print("ERROR")

def resu(frame ) :

     with open('Datadet.csv', 'r', encoding='utf-8') as csvfile:
                #sniff to find the format
                fileDialect = csv.Sniffer().sniff(csvfile.read(1024))
                csvfile.seek(0)
                #create a CSV reader
                myReader = csv.reader(csvfile, dialect=fileDialect)
                #read each row
                print(myReader)