#!/usr/bin/env python
from __future__ import division
import sys, os
from scripts.RunQC.Utilities import FileUtils as fileUtils
import xml.etree.ElementTree as ET
from scripts.RunQC.Metrics.TileMetrics import TileMetrics
from scripts.RunQC.Metrics.ErrorMetrics import ErrorMetrics

class MiSeqRunQualityInfo :

	# Checks phix alignment error rate
	def EvaluateErrorMetrics( self , config ) :

		path_element_list = config.basecalls_dir.split( 'Data' + os.path.sep + 'Intensities' + os.path.sep + 'BaseCalls')
		interop_path      = os.path.join( path_element_list[0] , 'InterOp' )
		error_metrics_file = os.path.join( interop_path , 'ErrorMetricsOut.bin' )

		if os.path.exists( error_metrics_file ) :

			metrics = ErrorMetrics( error_metrics_file )
			self.setAveragePhixErrorRate( metrics.average_error_rate )

	# calculates average percent failure amongst all tiles
	def CalculatePercentFailedDemux( self , config ) :

		alignment_dir = fileUtils.AlignmentDirectoryFromBaseCalls( config.basecalls_dir )

		demux_report  = os.path.join( alignment_dir , "DemultiplexSummaryF1L1.txt" )

		exists , message = fileUtils.exists( demux_report )

		if exists :

			percent_failed = list()

			foundHeader = False
			finished    = False
			for line in open( demux_report , 'rb' ) :
				if not foundHeader :
					if line.strip().startswith("SampleName") :
						foundHeader = True
						continue
					else : continue
				if foundHeader and not finished :
					tokens = line.strip().split('\t')
					if len( tokens ) > 2 :
						percent_failed.append( float( tokens[1] ) )
					else : finished = True

			if len( percent_failed ) != 0 :
				average_percent_failed = reduce( lambda x, y: x + y, percent_failed ) / len( percent_failed )
			else : average_percent_failed = 0

			self.setPercentFailedDemux( average_percent_failed )

	# Checks for percentage of tiles where zero reads demux.
	# Can indicate a problem with the flowcell - large spots, dust inhibiting image analysis or
	# it's also been seen when data is moved while it's being analyzed
	def CalculatePercentTileDemuxFail( self , config ) :

		alignment_dir = fileUtils.AlignmentDirectoryFromBaseCalls( config.basecalls_dir )

		demux_report  = os.path.join( alignment_dir , "DemultiplexSummaryF1L1.txt" )

		exists , message = fileUtils.exists( demux_report )

		if exists :

			zero_tiles  = 0
			total_tiles = 0

			foundHeader = False
			finished    = False
			for line in open( demux_report , 'rb' ) :
				if not foundHeader :
					if line.strip().startswith("SampleName") :
						foundHeader = True
						continue
					else : continue
				if foundHeader and not finished :
					tokens = line.strip().split('\t')
					if len( tokens ) > 2 :
						column_count = 0
						for token in tokens :
							if column_count > 1 :
								if float( token ) <= 0 :
									zero_tiles += 1
								total_tiles += 1
							column_count += 1
					else : finished = True

			if total_tiles > 0 :
				average_percent_failed = ( zero_tiles / total_tiles ) * 100
			else : average_percent_failed = 100

			self.setPercentTilesFailedDemux( average_percent_failed )

	# Checks for things like flow cell density etc.
	def EvaluateTileMetrics( self , config ) :

		path_element_list = config.basecalls_dir.split( 'Data' + os.path.sep + 'Intensities' + os.path.sep + 'BaseCalls')
		interop_path      = os.path.join( path_element_list[0] , 'InterOp' )

		tile_metrics_file = os.path.join( interop_path , 'TileMetricsOut.bin' )

		if os.path.exists( tile_metrics_file ) :

			metrics = TileMetrics( tile_metrics_file )

			self.setAverageDensity( metrics.average_density )
			self.setAverageDensityPF( metrics.average_density_pf )
			self.setPhasingApplied( metrics.average_phasing )
			self.setPrephasingApplied( metrics.average_prephasing )
			self.setPercentPhiXAligned( metrics.average_phix_aligned )
			self.setClusterCountPF( metrics.total_pf_clusters )
			self.setPercentPF( metrics.Get_Percent_PF() )

	# Sanity Check Fastq Files
	def VerifyFastqFilesExist( self , config ) :

		fastq_files_exist = True
		message = "File not found"

		for sampleNumber , sampleInfo in config.sampleInfoDictionary.iteritems() :
			fastq_files_exist , message = fileUtils.exists( sampleInfo.read_1_fastq )
			if not fastq_files_exist : break
			if config.is_paired_end :
				fastq_files_exist , message = fileUtils.exists( sampleInfo.read_2_fastq )
				if not fastq_files_exist : break

		self.setFastqFilesExist( fastq_files_exist , message )

	# Checks that the instrument completed
	# demux & fastq file generation
	def VerifyPrimaryAnalysisComplete( self , config ) :

		illumina_alignment_dir = fileUtils.AlignmentDirectoryFromBaseCalls( config.basecalls_dir )
		checkpoint_file        = os.path.join( illumina_alignment_dir, "Checkpoint.txt" )

		if os.path.exists( checkpoint_file ) :

			f = open( checkpoint_file , 'rb')
			line = f.readline().strip()
			f.close()

			step = int( line )

			if step >= 3 : self.setPrimaryAnalysisComplete( True , step )
			else : self.setPrimaryAnalysisComplete( False , step )

	# Warn if unknown illumina SW is detected.
	# sometimes they ship products that are not
	# ready for prime time - we should detect that
	def VerifyIlluminaSoftwareVersions( self , config ) :

		versions_file = os.path.join( os.path.dirname( __file__ ) , 'instrument_software_versions.xml' )

		if os.path.exists( versions_file  ) :

			known_versions_used = True

			tree = ET.parse( versions_file )
			root = tree.getroot()

			FPGA_versions = list()
			MCS_versions  = list()
			RTA_versions  = list()
			MSR_versions  = list()

			for version in root.findall( "./FPGA_versions/version" ) : # hardware
				FPGA_versions.append( version.text )
			for version in root.findall( "./MCS_versions/version" ) :  # firmware
				MCS_versions.append( version.text )
			for version in root.findall( "./RTA_versions/version" ) :  # real time anlaysis sw ( primary analysis )
				RTA_versions.append( version.text )
			for version in root.findall( "./MSR_versions/version" ) :  # miseq reporter ( secondary analysis )
				MSR_versions.append( version.text )

			versions = dict()

			versions['fpga_version'] = "unknown"
			versions['mcs_version']  = "unknown"
			versions['rta_version']  = "unknown"
			versions['msr_version']  = "unknown"

			# first 3 are read from run params file
			path_element_list = config.basecalls_dir.split( 'Data' + os.path.sep + 'Intensities' + os.path.sep + 'BaseCalls' )
			runParametersFile = os.path.join( path_element_list[0], 'RunParameters.xml' )
			if not os.path.exists( runParametersFile ) :
				runParametersFile = os.path.join( path_element_list[0], 'runParameters.xml' ) # sometimes it's lowercase!

			if os.path.exists( runParametersFile ) :

				tree = ET.parse( runParametersFile )
				root = tree.getroot()

				for child in root:
					if child.tag == "FPGAVersion" :
						versions['fpga_version'] = child.text
					elif child.tag == "MCSVersion" :
						versions['mcs_version'] = child.text
					elif child.tag == "RTAVersion" :
						versions['rta_version'] = child.text

			# Miseq reporter writes to generate fastq stats file
			generateFastqStatsFile = os.path.join( fileUtils.AlignmentDirectoryFromBaseCalls( config.basecalls_dir ) , 'GenerateFASTQRunStatistics.xml' )

			if os.path.exists( generateFastqStatsFile ) :

				tree = ET.parse( generateFastqStatsFile )
				root = tree.getroot()

				for child in root:
					if child.tag == "RunStats" :
						msr_element = child.find('AnalysisSoftwareVersion')
						versions['msr_version'] = msr_element.text

				if versions['fpga_version'] not in FPGA_versions :
					known_versions_used = False
				if versions['mcs_version'] not in MCS_versions :
					known_versions_used = False
				if versions['rta_version'] not in RTA_versions :
					known_versions_used = False
				if versions['msr_version'] not in MSR_versions :
					known_versions_used = False

				self.setMiSeqSWVersions( known_versions_used , versions )

	# Given a base call dir, write the report
	def writeReport( self , config , version, working_dir ) :

		#file = open( os.path.join( config.run_root , "RunQuality_" + config.flowcellID + ".csv" ) , 'wb' )
		if not os.path.exists( os.path.join( working_dir , config.runID ) ):
			os.makedirs( os.path.join( working_dir , config.runID ) )
		file = open( os.path.join( working_dir , config.runID, config.runID + "_stats.csv" ) , 'wb' )

		file.write( "FLAG,METRIC,QUALITY_SCORE,MESSAGE,VERSION,RUN_ID\n" )
		#file.write( "Instrument Software Check" +
				#',' + 'N/A' +
				#',' + str( self.known_sw_versions_used_score ) +
				#',' + self.known_sw_versions_used_msg +
				#',' + version +
				#',' + config.flowcellID + '\n' )
		#file.write( "Fastq File Sanity Check" +
				#',' + 'N/A' +
				#',' + str( self.fastq_files_score ) +
				#',' + self.fastq_files_exist_msg +
				#',' + version +
				#',' + config.flowcellID + '\n' )
		#file.write( "Instrument Run Complete" +
				#',' + str( self.last_primary_analysis_step ) +
				#',' + str( self.primary_analysis_complete_score ) +
				#',' + self.primary_analysis_complete_msg +
				#',' + version +
				#',' + config.flowcellID + '\n' )
		file.write( "Phasing" +
				',' + str( "%0.2f" % self.average_phasing ) +
				',' + str( self.phasing_score )  +
				',' + self.phasing_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Prephasing" +
				',' + str( "%0.2f" % self.average_prephasing ) +
				',' + str( self.prephasing_score )  +
				',' + self.prephasing_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "PhiX Control Error Rate" +
				',' + str( "%0.2f" % self.average_phix_error_rate ) +
				',' + str( self.average_phix_error_rate_score )  +
				',' + self.average_phix_error_rate_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Percent PF"
				',' + str( "%0.2f" % self.percent_pf ) +
				',' + str( self.percent_pf_score ) +
				',' + self.percent_pf_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Cluster Density" +
				',' + str( int( self.cluster_density ) ) +
				',' + str( self.cluster_density_score ) +
				',' + self.cluster_density_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Cluster Density PF" +
				',' + str( int( self.cluster_density_pf ) ) +
				',' + str( self.cluster_density_pf_score ) +
				',' + self.cluster_density_pf_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Reads PF" +
				',' + str( self.total_pf_clusters ) +
				',' + str( self.reads_pf_score ) +
				',' + self.reads_pf_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Percent PhiX Aligned" +
				',' + str( "%0.2f" % self.average_phix_aligned ) +
				',' + str( self.percent_phix_aligned_score ) +
				',' + self.percent_phix_aligned_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Percent Failed Demux" +
				',' + str( "%0.2f" % self.average_demux_failed ) +
				',' + str( self.percent_demux_failed_score ) +
				',' + self.percent_demux_failed_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.write( "Percent Tiles Failed Demux" +
				',' + str( "%0.2f" % self.percent_tile_demux_fail ) +
				',' + str( self.percent_tile_demux_fail_score ) +
				',' + self.percent_tile_demux_fail_msg +
				',' + version +
				',' + config.irunID + '\n' )
		file.close()

	# Given a base call dir, write the report
	def writeReport2( self , config , version, working_dir ) :

		#file = open( os.path.join( config.run_root , "RunQuality_" + config.flowcellID + ".csv" ) , 'wb' )
		if not os.path.exists( os.path.join( working_dir , config.runID ) ):
			os.makedirs( os.path.join( working_dir , config.runID ) )
		file = open( os.path.join( working_dir , config.runID, config.runID + "_stats.csv" ) , 'ab' )

		file.write( "RunId,Tool,FLAG,StandardGroup,Type,Metric,Comments\n" )
		wantedMetricsNames = ['Phasing','Prephasing','PhiX Control Error Rate','Percent PF','Cluster Density','Cluster Density PF','Reads PF','Percent PhiX Aligned','Percent Failed Demux','Percent Tiles Failed Demux']
		wantedMetricsVariables = ['average_phasing','average_prephasing','average_phix_error_rate','percent_pf','cluster_density','cluster_density_pf','total_pf_clusters','average_phix_aligned','average_demux_failed','percent_tile_demux_fail']
		wantedMetricsMsgs = ['phasing_msg','prephasing_msg','average_phix_error_rate_msg','percent_pf_msg','cluster_density_msg','cluster_density_pf_msg','reads_pf_msg','percent_phix_aligned_msg','percent_demux_failed_msg','percent_tile_demux_fail_msg']

		if len(wantedMetricsNames) == len(wantedMetricsVariables) == len(wantedMetricsMsgs)  :
			for index,MetricName in enumerate(wantedMetricsNames):
				MetricVariable = wantedMetricsVariables[index]
				MetricMsg = wantedMetricsMsgs[index]

				file.write( config.irunID +
					',' + version +
					',' + MetricName +
					',' + 'ALL' +
					',' + 'ALL' +
					',' + str( "%0.2f" % self.__dict__[MetricVariable] ) +
					',' + self.__dict__[MetricMsg] + '\n'
				)

		file.close()

	# Average error rate from RTA SeqAn alignment of phix
	def setAveragePhixErrorRate( self , average_error_rate ) :

		self.average_phix_error_rate = average_error_rate

		if average_error_rate <= self.max_error_rate :
			self.average_phix_error_rate_msg = "PhiX Error Rate [spec: <= " + str( "%.2f" % self.max_error_rate ) + "]: " + str( "%.2f" % average_error_rate )
			self.average_phix_error_rate_ok = True
			self.average_phix_error_rate_score = 10

		else :
			self.average_phix_error_rate_ok = True
			self.average_phix_error_rate_msg = "PhiX Error Rate is greater than spec [spec: <= " + str( "%.2f" % self.max_error_rate ) + "]: " + str( "%.2f" % average_error_rate )

			if average_error_rate > self.max_error_rate and average_error_rate <= self.max_error_rate + 0.1 : self.average_phix_error_rate_score = 9
			elif average_error_rate > self.max_error_rate + 0.1 and average_error_rate <= self.max_error_rate + 0.2 : self.average_phix_error_rate_score = 8
			elif average_error_rate > self.max_error_rate + 0.2 and average_error_rate <= self.max_error_rate + 0.3 : self.average_phix_error_rate_score = 7
			elif average_error_rate > self.max_error_rate + 0.3 and average_error_rate <= self.max_error_rate + 0.4 : self.average_phix_error_rate_score = 6
			elif average_error_rate > self.max_error_rate + 0.4 and average_error_rate <= self.max_error_rate + 0.5 : self.average_phix_error_rate_score = 5
			elif average_error_rate > self.max_error_rate + 0.5 and average_error_rate <= self.max_error_rate + 0.6 : self.average_phix_error_rate_score = 4
			elif average_error_rate > self.max_error_rate + 0.6 and average_error_rate <= self.max_error_rate + 0.7 : self.average_phix_error_rate_score = 3
			elif average_error_rate > self.max_error_rate + 0.7 and average_error_rate <= self.max_error_rate + 0.8 : self.average_phix_error_rate_score = 2
			elif average_error_rate > self.max_error_rate + 0.8 : self.average_phix_error_rate_score = 1 ; self.average_phix_error_rate_ok = False

	# Average demux failure percent amongst all tiles
	def setPercentFailedDemux( self , average_percent_failed ) :

		self.average_demux_failed = average_percent_failed

		if average_percent_failed <= self.max_mean_demux_failure :

			self.percent_demux_failed_msg = "Average percent demux failure is within spec [spec: <= " + str( self.max_mean_demux_failure ) + "]: " + str( "%.2f" % average_percent_failed )
			self.percent_demux_failed_ok = True
			self.percent_demux_failed_score = 10

		else :
			self.percent_demux_failed_ok = True
			self.percent_demux_failed_msg = "Average percent demux failure is greater than spec [spec: <= " + str( self.max_mean_demux_failure ) + "]: " + str( "%.2f" % average_percent_failed)
			self.percent_demux_failed_score = 10

			if average_percent_failed > self.max_mean_demux_failure         and average_percent_failed < self.max_mean_demux_failure + 3  : self.percent_demux_failed_score = 9
			elif average_percent_failed >= self.max_mean_demux_failure + 3  and average_percent_failed < self.max_mean_demux_failure + 7  : self.percent_demux_failed_score = 8
			elif average_percent_failed >= self.max_mean_demux_failure + 7  and average_percent_failed < self.max_mean_demux_failure  + 11 : self.percent_demux_failed_score = 7
			elif average_percent_failed >= self.max_mean_demux_failure + 11 and average_percent_failed < self.max_mean_demux_failure + 14 : self.percent_demux_failed_score = 6
			elif average_percent_failed >= self.max_mean_demux_failure + 14 and average_percent_failed < self.max_mean_demux_failure + 16 : self.percent_demux_failed_score = 5
			elif average_percent_failed >= self.max_mean_demux_failure + 16 and average_percent_failed < self.max_mean_demux_failure + 18 : self.percent_demux_failed_score = 4
			elif average_percent_failed >= self.max_mean_demux_failure + 18 and average_percent_failed < self.max_mean_demux_failure + 22 : self.percent_demux_failed_score = 3
			elif average_percent_failed >= self.max_mean_demux_failure + 22 and average_percent_failed < self.max_mean_demux_failure + 25 : self.percent_demux_failed_score = 2
			elif average_percent_failed >= self.max_mean_demux_failure + 25 : self.percent_demux_failed_score = 1 ; self.percent_demux_failed_ok = False

	# Percentage of tiles where demux completely failed.  Indicates problem
	# with the flowcell like dust.  Can also be caused by moving the data while
	# it's being analyzed - the threads die, the fastq 'appears' valid ( not truncated )
	# but it's missing reads..
	def setPercentTilesFailedDemux( self , percent_failed ) :

		self.percent_tile_demux_fail_msg = percent_failed

		if percent_failed <= self.max_percent_tile_demux_fail :

			self.percent_tile_demux_fail_msg   = "Average percent demux tile failure is within spec [spec: <= " + str( self.max_percent_tile_demux_fail ) + "]: " + str( "%.2f" % percent_failed )
			self.percent_tile_demux_fail_ok    = True
			self.percent_tile_demux_fail_score = 10
			self.percent_tile_demux_fail       = percent_failed

		else :

			self.percent_tile_demux_fail_msg   = "Average percent demux tile failure is greater than spec [spec: <= " + str( self.max_percent_tile_demux_fail ) + "]: " + str( "%.2f" % percent_failed )
			self.percent_tile_demux_fail_ok    = False
			self.percent_tile_demux_fail_score = 10
			self.percent_tile_demux_fail       = percent_failed

			if percent_failed > self.max_percent_tile_demux_fail         and percent_failed < self.max_percent_tile_demux_fail + 3  : self.percent_tile_demux_fail_score = 9
			elif percent_failed >= self.max_percent_tile_demux_fail + 3  and percent_failed < self.max_percent_tile_demux_fail + 7  : self.percent_tile_demux_fail_score = 8
			elif percent_failed >= self.max_percent_tile_demux_fail + 7  and percent_failed < self.max_percent_tile_demux_fail + 11 : self.percent_tile_demux_fail_score = 7
			elif percent_failed >= self.max_percent_tile_demux_fail + 11 and percent_failed < self.max_percent_tile_demux_fail + 14 : self.percent_tile_demux_fail_score = 6
			elif percent_failed >= self.max_percent_tile_demux_fail + 14 and percent_failed < self.max_percent_tile_demux_fail + 16 : self.percent_tile_demux_fail_score = 5
			elif percent_failed >= self.max_percent_tile_demux_fail + 16 and percent_failed < self.max_percent_tile_demux_fail + 18 : self.percent_tile_demux_fail_score = 4
			elif percent_failed >= self.max_percent_tile_demux_fail + 18 and percent_failed < self.max_percent_tile_demux_fail + 22 : self.percent_tile_demux_fail_score = 3
			elif percent_failed >= self.max_percent_tile_demux_fail + 22 and percent_failed < self.max_percent_tile_demux_fail + 25 : self.percent_tile_demux_fail_score = 2
			elif percent_failed >= self.max_percent_tile_demux_fail + 25 : self.percent_tile_demux_fail_score = 1

	# Sometimes illumina ships half-baked and if we encounter
	# an unknown we want to warn at least
	def setMiSeqSWVersions( self , known_sw_versions_used , versions ) :

		if known_sw_versions_used :
			self.known_sw_versions_used_msg = "Known & tested instrument SW was used to generate fastq files"
			self.known_sw_versions_used_ok = True
			self.known_sw_versions_used_score = 10
		else :
			self.known_sw_versions_used_msg = "WARNING: Unknown & untested (by Trovagene) instrument SW was used to generate data!"
			self.known_sw_versions_used_ok = False
			self.known_sw_versions_used_score = 1

	# Checks that the instrument dropped the flag
	# indicating step-3 completed ( finished making fastqs )
	def setPrimaryAnalysisComplete( self , complete , step ) :

		if complete == True :
			self.primary_analysis_complete_msg = "Min required step: 3 achieved: 3"
			self.primary_analysis_complete_ok = True
			self.primary_analysis_complete_score = 10
			self.last_primary_analysis_step = step
		else :
			self.primary_analysis_complete_msg = "Min required step (3) not achieved: " + str( complete )
			self.primary_analysis_complete_ok = False
			self.primary_analysis_complete_score = -1
			self.last_primary_analysis_step = step

	# Sanity check to verify fastq files exist
	def setFastqFilesExist( self , exist , msg ) :

		if exist :
			self.fastq_files_exist_ok = True
			self.fastq_files_exist_msg = "All Fastq files exist"
			self.fastq_files_score = 10
		else :
			self.fastq_files_exist_ok = False
			self.fastq_files_exist_msg = msg
			self.fastq_files_score = -1

	# The minimum required number of clusters passed the chastity filter:
	# highest_sig_int / (highest_sig_int + 2nd_highest_sig_int)
	# Clusters pass filter if no more than one base call in the first 25 cycles has a chastity of less than 0.6
	def setPercentPF( self , percent_pf ) :

		self.percent_pf = percent_pf

		if percent_pf >= self.min_percent_PF :
			self.percent_pf_msg = "Percent PF is within spec [spec: >= " + str( self.min_percent_PF ) + "]: " + str( "%.2f" % percent_pf )
			self.percent_pf_ok = True
			self.percent_pf_score = 10

		else : # somewhat sigmoidal drop is qualites as pf decreases

			self.percent_pf_msg = "Percent PF is less than spec [spec: >= " + str( self.min_percent_PF ) + "]: " + str( "%.2f" % percent_pf )

			if percent_pf   < self.min_percent_PF - 0.05 : self.percent_pf_score = 9
			elif percent_pf < self.min_percent_PF - 0.09 : self.percent_pf_score = 8
			elif percent_pf < self.min_percent_PF - 0.13 : self.percent_pf_score = 7
			elif percent_pf < self.min_percent_PF - 0.15 : self.percent_pf_score = 6
			elif percent_pf < self.min_percent_PF - 0.17 : self.percent_pf_score = 5
			elif percent_pf < self.min_percent_PF - 0.19 : self.percent_pf_score = 4
			elif percent_pf < self.min_percent_PF - 0.20 : self.percent_pf_score = 3
			elif percent_pf < self.min_percent_PF - 0.21 : self.percent_pf_score = 2
			elif percent_pf < self.min_percent_PF - 0.22 : self.percent_pf_score = 1

	# The PF cluster count should be within the density specification set by illumina for the V3 chemistry
    # Single Reads: 22-25 M
    # Paired End Reads: 44-50 M
    # For now while we are sequencing mono-template we may actually get better performance with lower density.
    # Perhaps even to the point of obtaining a very low quality score using the scale implemented.
    # A test we should add is a measure of the distance between Raw clusters and PF clusters.  If things went
    # well they should be close together - but how close?
	def setClusterCountPF( self , total_pf_clusters ) :

		self.total_pf_clusters = total_pf_clusters

		if total_pf_clusters > self.min_reads_PF and total_pf_clusters < self.max_reads_PF :

			self.reads_pf_msg = "The expected number of reads are available: " + str( "%.0f" % total_pf_clusters )
			self.reads_pf_ok = True
			self.reads_pf_score = 10

		elif total_pf_clusters > self.max_reads_PF :

				self.reads_pf_msg = "The number of reads passing filter exceeds the spec (" + str( "%.0f" % self.max_reads_PF ) + "): " + str( "%.0f" % total_pf_clusters )
				self.reads_pf_ok = True

				if total_pf_clusters    > self.max_reads_PF + 1000000 and total_pf_clusters < self.max_reads_PF + 1500000 : self.reads_pf_score = 9
				elif total_pf_clusters >= self.max_reads_PF + 1500000 and total_pf_clusters < self.max_reads_PF + 2000000 : self.reads_pf_score = 8
				elif total_pf_clusters >= self.max_reads_PF + 2000000 and total_pf_clusters < self.max_reads_PF + 2500000 : self.reads_pf_score = 7
				elif total_pf_clusters >= self.max_reads_PF + 2500000 and total_pf_clusters < self.max_reads_PF + 3000000 : self.reads_pf_score = 6
				elif total_pf_clusters >= self.max_reads_PF + 3000000 and total_pf_clusters < self.max_reads_PF + 3500000 : self.reads_pf_score = 5
				elif total_pf_clusters >= self.max_reads_PF + 3500000 and total_pf_clusters < self.max_reads_PF + 4000000 : self.reads_pf_score = 4
				elif total_pf_clusters >= self.max_reads_PF + 4000000 and total_pf_clusters < self.max_reads_PF + 4500000 : self.reads_pf_score = 3
				elif total_pf_clusters >= self.max_reads_PF + 4500000 and total_pf_clusters < self.max_reads_PF + 5000000 : self.reads_pf_score = 2
				elif total_pf_clusters >= self.max_reads_PF + 5000000                                                     : self.reads_pf_score = 1

		elif total_pf_clusters < self.min_reads_PF :

				self.reads_pf_msg = "The number of reads passing filter does not meet spec (" + str( "%.0f" % self.min_reads_PF ) + "): " + str( "%.0f" % total_pf_clusters )

				if   total_pf_clusters > self.min_reads_PF - 2000000  : self.reads_pf_score = 9 ; self.reads_pf_ok = True
				elif total_pf_clusters > self.min_reads_PF - 4000000  : self.reads_pf_score = 8 ; self.reads_pf_ok = True
				elif total_pf_clusters > self.min_reads_PF - 8000000  : self.reads_pf_score = 7 ; self.reads_pf_ok = True
				elif total_pf_clusters > self.min_reads_PF - 10000000 : self.reads_pf_score = 6 ; self.reads_pf_ok = True
				elif total_pf_clusters > self.min_reads_PF - 12000000 : self.reads_pf_score = 5
				elif total_pf_clusters > self.min_reads_PF - 14000000 : self.reads_pf_score = 4
				elif total_pf_clusters > self.min_reads_PF - 16000000 : self.reads_pf_score = 3
				elif total_pf_clusters > self.min_reads_PF - 17000000 : self.reads_pf_score = 2
				elif total_pf_clusters > 9000000                      : self.reads_pf_score = 1

	# Sets the percentage of PhiX which aligned.  Currently we are spiking in 5% and so we'd expect to see 5% aligned.
	# If not it may be due to:
	# 1. PhiX sample degradation - a tube freeze/thawed multiple times over a week of use may degrade so that even though
	#    the lab thinks they are spiking in 5% it is now 2%. Recommend aliquot and store single use tubes.
	# 2. Can also indicate overall poor sequencing quality if the error (1) is avoided.
	def setPercentPhiXAligned( self , average_phix_aligned ) :

		self.average_phix_aligned = average_phix_aligned

		if average_phix_aligned >= self.min_percent_phix_aligned :

			self.percent_phix_aligned_msg = "Percent PhiX Aligned is within spec [spec: >= " + \
			str( "%.2f" % self.min_percent_phix_aligned ) + "]: " + str( "%.2f" % average_phix_aligned )

			# Too much PHIX indicates low to no human
			if average_phix_aligned < self.max_percent_phix_aligned :
				self.percent_phix_aligned_ok = True
				self.percent_phix_aligned_score = 10
			else :
				self.percent_phix_aligned_ok = False
				self.percent_phix_aligned_score = 1
				self.percent_phix_aligned_msg = "Percent PhiX Aligned indicates overloading PhiX or underloading Human: " + str( "%.2f" % average_phix_aligned )
		else :

			self.percent_phix_aligned_msg = "Percent PhiX Aligned is less than spec [spec: " + \
			str( "%.2f" %  self.min_percent_phix_aligned ) + "]: " + str( "%.2f" % average_phix_aligned )

			if   average_phix_aligned < self.min_percent_phix_aligned and average_phix_aligned >= self.min_percent_phix_aligned - 0.5 :       self.percent_phix_aligned_score = 9 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 0.5 and average_phix_aligned >= self.min_percent_phix_aligned - 1 :   self.percent_phix_aligned_score = 8 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 1   and average_phix_aligned >= self.min_percent_phix_aligned - 1.5 : self.percent_phix_aligned_score = 7 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 1.5 and average_phix_aligned >= self.min_percent_phix_aligned - 2 :   self.percent_phix_aligned_score = 6 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 2   and average_phix_aligned >= self.min_percent_phix_aligned - 2.5 : self.percent_phix_aligned_score = 5 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 2.5 and average_phix_aligned >= self.min_percent_phix_aligned - 3 :   self.percent_phix_aligned_score = 4 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 3   and average_phix_aligned >= self.min_percent_phix_aligned - 3.5 : self.percent_phix_aligned_score = 3 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 3.5 and average_phix_aligned >= self.min_percent_phix_aligned - 4 :   self.percent_phix_aligned_score = 2 ; self.percent_phix_aligned_ok = True
			elif average_phix_aligned < self.min_percent_phix_aligned - 4 :                                                                   self.percent_phix_aligned_score = 1 ; self.percent_phix_aligned_ok = True

	# Sets the Prephasing applied status: Correction for % of DNA strands per cluster which
	# have erroneously fallen ahead of the current cycle by one or more bases
	def setPrephasingApplied( self , average_prephasing ) :

		self.average_prephasing = average_prephasing

		if average_prephasing <= self.max_prephasing :

			self.prephasing_score = 10
			self.prephasing_msg = "Prephasing applied is within spec [spec: <= " + str( "%.2f" % self.max_prephasing ) + "]: " + str( "%.2f" % average_prephasing )
			self.prephasing_ok = True

		else : # primitive quality score for now...

			difference = average_prephasing - self.max_prephasing

			if difference < 0.1 : self.prephasing_score = 9
			elif difference >= 0.1 and difference <= 0.2 : self.prephasing_score = 8
			elif difference > 0.2 and difference <= 0.3 : self.prephasing_score = 7
			elif difference > 0.3 and difference <= 0.4 : self.prephasing_score = 6
			elif difference > 0.4 and difference <= 0.5 : self.prephasing_score = 5
			elif difference > 0.5 and difference <= 0.6 : self.prephasing_score = 4
			elif difference > 0.6 and difference <= 0.7 : self.prephasing_score = 3
			elif difference > 0.7 and difference <= 0.8 : self.prephasing_score = 2
			elif difference > 0.8 : self.prephasing_score = 1

			self.prephasing_msg = "Prephasing applied is greater than spec [spec: <= " + str( "%.2f" % self.max_prephasing ) + "]: " + \
			str( "%.2f" % average_prephasing ) + ". Difference is: " + str( "%.2f" % difference )

	# Sets the Phasing applied status: Correction for % of DNA strands per cluster which
	# have erroneously fallen behind the current cycle by one or more bases
	def setPhasingApplied( self , average_phasing ) :

		self.average_phasing = average_phasing

		if average_phasing <= self.max_phasing :

			self.phasing_score = 10
			self.phasing_msg = "Phasing applied is within spec [spec: <= " + str( "%.2f" % self.max_phasing ) + "]: " + str( "%.2f" % average_phasing )
			self.phasing_ok = True

		else : # primitive quality score for now...

			difference = average_phasing - self.max_phasing

			if difference < 0.1 : self.phasing_score = 9
			elif difference >= 0.1 and difference <= 0.2 : self.phasing_score = 8
			elif difference > 0.2 and difference <= 0.3 : self.phasing_score = 7
			elif difference > 0.3 and difference <= 0.4 : self.phasing_score = 6
			elif difference > 0.4 and difference <= 0.5 : self.phasing_score = 5
			elif difference > 0.5 and difference <= 0.6 : self.phasing_score = 4
			elif difference > 0.6 and difference <= 0.7 : self.phasing_score = 3
			elif difference > 0.7 and difference <= 0.8 : self.phasing_score = 2
			elif difference > 0.8 : self.phasing_score = 1

			self.phasing_msg = "Phasing applied is greater than spec [spec: <= " + str( "%.2f" % self.max_phasing ) + "]: " + \
			str( "%.2f" % average_phasing ) + ". Difference is: " + str( "%.2f" % difference )

	# One of the most important metrics relating run quality.
	# More important that it's not too dense (over-cooked) rather than not dense enough (undercooked)
	# For now while we are sequencing mono-template we may actually get better performance with lower density.
	# Perhaps even to the point of obtaining a very low quality score using the scale implemented.
    # A test we should add is a measure of the distance between Raw clusters and PF clusters.  If things went
    # well they should be close together - but how close?
	def setAverageDensity( self , average_density ) :

		self.cluster_density = average_density

		if average_density <= self.max_cluster_density and average_density >= self.min_cluster_density :

			self.cluster_density_msg = "Raw cluster density is within spec [spec: >= " + str( int( self.min_cluster_density / 1000 ) ) + "K/mm^2 and <= " + \
			str( int( self.max_cluster_density / 1000 ) ) + "K/mm^2]: " + str( int( average_density / 1000 ) ) + "K/mm^2"

			self.cluster_density_ok = True
			self.cluster_density_score = 10

		else :

			if average_density < self.min_cluster_density : # undercooked, better than overcooked
				percent_difference = abs( self.min_cluster_density - average_density ) / abs( ( self.min_cluster_density + average_density ) / 2 ) * 100
				if percent_difference < 15 : self.cluster_density_score = 9
				elif percent_difference >= 15 and percent_difference < 35   : self.cluster_density_score = 8
				elif percent_difference >= 35 and percent_difference < 50   : self.cluster_density_score = 7
				elif percent_difference >= 50 and percent_difference < 65   : self.cluster_density_score = 6
				elif percent_difference >= 65 and percent_difference < 80   : self.cluster_density_score = 5
				elif percent_difference >= 80 and percent_difference < 95   : self.cluster_density_score = 4
				elif percent_difference >= 95 and percent_difference < 110  : self.cluster_density_score = 3
				elif percent_difference >= 110 and percent_difference < 125 : self.cluster_density_score = 2
				elif percent_difference >= 125 : self.cluster_density_score = 1

			elif average_density > self.max_cluster_density :
				percent_difference = abs( self.max_cluster_density - average_density ) / abs( ( self.max_cluster_density + average_density ) / 2 ) * 100
				if percent_difference < 1 : self.cluster_density_score = 9
				elif percent_difference >= 1 and percent_difference < 5   : self.cluster_density_score = 8
				elif percent_difference >= 5 and percent_difference < 10  : self.cluster_density_score = 7
				elif percent_difference >= 10 and percent_difference < 15 : self.cluster_density_score = 6
				elif percent_difference >= 15 and percent_difference < 20 : self.cluster_density_score = 5
				elif percent_difference >= 20 and percent_difference < 25 : self.cluster_density_score = 4
				elif percent_difference >= 25 and percent_difference < 30 : self.cluster_density_score = 3
				elif percent_difference >= 30 and percent_difference < 35 : self.cluster_density_score = 2
				elif percent_difference >= 35 : self.cluster_density_score = 1

			self.cluster_density_msg = "Raw cluster density is not within spec [spec: >= " + str( "%.0f" % ( self.min_cluster_density / 1000 ) ) + "K/mm^2 and <= " + \
                str( "%.0f" % ( self.max_cluster_density  / 1000 ) ) + "K/mm^2]: " + str( "%.0f" % ( average_density / 1000 ) )  + \
				" K/mm^2. Percent Difference is: " + str( "%.0f" % ( percent_difference ) )

		# This is density for data which has passed the chastity filter.  The raw density
		# is more informative but it's not always estimated correctly by the illumina
		# software.  PF density should be lower but relatively close if the true
		# density is in the sweet spot of the illumina image analysis algorithms.
		# For now while we are sequencing mono-template we may actually get better performance with lower density.
		# Perhaps even to the point of obtaining a very low quality score using the scale implemented.
		# A test we should add is a measure of the distance between Raw clusters and PF clusters.  If things went
		# well they should be close together - but how close?
	def setAverageDensityPF( self , average_density_pf ) :

		self.cluster_density_pf = average_density_pf

		if average_density_pf <= self.max_cluster_density and average_density_pf >= self.min_cluster_density :

			self.cluster_density_pf_msg = "PF cluster density is within spec [spec: >= " + str( int( self.min_cluster_density / 1000 ) ) + "K/mm^2 and <= " + \
			str( int( self.max_cluster_density / 1000 ) ) + "K/mm^2]: " + str( int( average_density_pf / 1000 ) ) + "K/mm^2"

			self.cluster_density_pf_ok = True
			self.cluster_density_pf_score = 10

		else :

			if average_density_pf < self.min_cluster_density : # undercooked, better than overcooked
				percent_difference = abs( self.min_cluster_density - average_density_pf ) / abs( ( self.min_cluster_density + average_density_pf ) / 2 ) * 100
				if percent_difference < 15 : self.cluster_density_pf_score = 9
				elif percent_difference >= 15 and percent_difference < 35   : self.cluster_density_pf_score = 8
				elif percent_difference >= 35 and percent_difference < 50   : self.cluster_density_pf_score = 7
				elif percent_difference >= 50 and percent_difference < 65   : self.cluster_density_pf_score = 6
				elif percent_difference >= 65 and percent_difference < 80   : self.cluster_density_pf_score = 5
				elif percent_difference >= 80 and percent_difference < 95   : self.cluster_density_pf_score = 4
				elif percent_difference >= 95 and percent_difference < 110  : self.cluster_density_pf_score = 3
				elif percent_difference >= 110 and percent_difference < 125 : self.cluster_density_pf_score = 2
				elif percent_difference >= 125 : self.cluster_density_pf_score = 1

			elif average_density_pf > self.max_cluster_density :
				percent_difference = abs( self.max_cluster_density - average_density_pf ) / abs( ( self.max_cluster_density + average_density_pf ) / 2 ) * 100
				if percent_difference < 1 : self.cluster_density_pf_score = 9
				elif percent_difference >= 1 and percent_difference < 5   : self.cluster_density_pf_score = 8
				elif percent_difference >= 5 and percent_difference < 10  : self.cluster_density_pf_score = 7
				elif percent_difference >= 10 and percent_difference < 15 : self.cluster_density_pf_score = 6
				elif percent_difference >= 15 and percent_difference < 20 : self.cluster_density_pf_score = 5
				elif percent_difference >= 20 and percent_difference < 25 : self.cluster_density_pf_score = 4
				elif percent_difference >= 25 and percent_difference < 30 : self.cluster_density_pf_score = 3
				elif percent_difference >= 30 and percent_difference < 35 : self.cluster_density_pf_score = 2
				elif percent_difference >= 35 : self.cluster_density_pf_score = 1

			self.cluster_density_pf_msg = "PF cluster density is not within spec [spec: >= " + str( "%.0f" % ( self.min_cluster_density / 1000 ) ) + "K/mm^2 and <= " + \
                str( "%.0f" % ( self.max_cluster_density / 1000 ) ) + "K/mm^2]: " + str( "%.0f" % ( average_density_pf / 1000 ) )  + \
				" K/mm^2. Percent Difference is: " + str( "%.0f" % ( percent_difference ) )

	# Failure of any of these results in a "No" answer to the
	# question: "Is this data good enough for teriary analysis?"
	# Secondary analysis being fastq file generation via MiSeq Reporter
	def OK( self ) :
		if self.fastq_files_exist_ok and self.primary_analysis_complete_ok and self.average_phix_error_rate_ok and \
		   self.known_sw_versions_used_ok and self.reads_pf_ok and self.percent_phix_aligned_ok and \
		   self.percent_demux_failed_ok and self.percent_tile_demux_fail_ok : return True , "AOK"
		elif not self.fastq_files_exist_ok :
			return False , self.fastq_files_exist_msg
		elif not self.primary_analysis_complete_ok :
			return False , self.primary_analysis_complete_msg
		elif not self.average_phix_error_rate_ok :
			return False , self.average_phix_error_rate_msg
		elif not self.known_sw_versions_used_ok :
			return False , self.known_sw_versions_used_msg
		elif not self.reads_pf_ok :
			return False , self.reads_pf_msg
		elif not self.percent_phix_aligned_ok :
			return False , self.percent_phix_aligned_msg
		elif not self.percent_demux_failed_ok :
			return False , self.percent_demux_failed_msg
		elif not self.percent_tile_demux_fail_ok :
			return False , self.percent_tile_demux_fail_msg

	####################################
	# MAIN INITIALIZATION

	def __init__( self , config ):

		self.average_phix_aligned    = 0
		self.total_pf_clusters       = 0
		self.percent_pf              = 0
		self.average_prephasing      = 0
		self.average_phasing         = 0
		self.cluster_density         = 0
		self.cluster_density_pf      = 0
		self.average_demux_failed    = 0
		self.average_phix_error_rate = 0
		self.percent_tile_demux_fail = 0

		self.cluster_density_msg           = "File not found"
		self.cluster_density_pf_msg        = "File not found"
		self.phasing_msg                   = "File not found"
		self.prephasing_msg                = "File not found"
		self.percent_phix_aligned_msg      = "File not found"
		self.cluster_count_pf_msg          = "File not found"
		self.percent_pf_msg                = "File not found"
		self.reads_pf_msg                  = "File not found"
		self.fastq_files_exist_msg         = "File not found"
		self.primary_analysis_complete_msg = "File not found"
		self.known_sw_versions_used_msg    = "File not found"
		self.percent_demux_failed_msg      = "File not found"
		self.average_phix_error_rate_msg   = "File not found"
		self.percent_tile_demux_fail_msg   = "File not found"

		self.cluster_density_ok           = False
		self.cluster_density_pf_ok        = False
		self.phasing_ok                   = False
		self.prephasing_ok                = False
		self.percent_phix_aligned_ok      = False
		self.reads_pf_ok                  = False
		self.percent_pf_ok                = False
		self.fastq_files_exist_ok         = False
		self.primary_analysis_complete_ok = False
		self.known_sw_versions_used_ok    = False
		self.percent_demux_failed_ok      = False
		self.average_phix_error_rate_ok   = False
		self.percent_tile_demux_fail_ok   = False

		self.cluster_density_score           = 0
		self.cluster_density_pf_score        = 0
		self.phasing_score                   = 0
		self.prephasing_score                = 0
		self.percent_phix_aligned_score      = 0
		self.reads_pf_score                  = 0
		self.percent_pf_score                = 0
		self.fastq_files_score               = 0
		self.primary_analysis_complete_score = 0
		self.known_sw_versions_used_score    = 0
		self.percent_demux_failed_score      = 0
		self.average_phix_error_rate_score   = 0
		self.percent_tile_demux_fail_score   = 0

		self.last_primary_analysis_step = 3    # an illumina workflow step, last step in GenerateFASTQ workflow
		self.min_reads_PF = 17500000           # based on cluster density for v3 chemistry:(1200-1400 K/m) diverse or 500-1200 K/m monotemplate PF.  For diverse: 22000000
		self.max_reads_PF = 22000000           # based on cluster density for v3 chemsitry: (1200-1400 K/m diverse or 500-1200 K/m monotemplate PF ) For diverse: 25000000
		self.min_percent_phix_aligned = 5      # currently spiking in 5% but more is better for mono-template. Lower score if degrades.
		self.max_percent_phix_aligned = 60     # indicates at most only 40% of the flowcell is human, which is not good
		self.min_percent_PF = 75               # not too conservative considering monotemplate
		self.max_error_rate = 0.50             # Control PhiX RTA/SeqAn alignment error rate
		self.max_phasing = 0.5                 # percent of clusters falling behind (ideal is < 0.25)
		self.max_prephasing = 0.3              # percent clusters getting ahead (ideal is < 0.15)
		self.max_cluster_density = 1200000     # MiSeq v3 reagents accommodate an optimal raw cluster density of ( 1200-1400 K/m diverse or 500-1200 K/m monotemplate )
		self.min_cluster_density = 500000      # MiSeq v3 reagents accommodate an optimal raw cluster density of ( 1200-1400 K/m diverse or 500-1200 K/m monotemplate )
		self.max_mean_demux_failure = 20       # 5-10% is considered good, this is liberal because we're sequencing mono-template
		self.max_percent_tile_demux_fail = 10  # Percentage of tiles were zero reads demultiplex

		self.VerifyFastqFilesExist( config )
		self.VerifyPrimaryAnalysisComplete( config )
		self.VerifyIlluminaSoftwareVersions( config )
		self.EvaluateTileMetrics( config )
		self.CalculatePercentFailedDemux( config )
		self.CalculatePercentTileDemuxFail( config )
		self.EvaluateErrorMetrics( config )


# -*- coding: utf-8 -*-
import torch.nn as nn
import math

import torch
import torch.nn.functional as F

import numpy as np
import copy

from utils.arch_util import ESA, PA

from RDU import DynamicDiverseBranchBlock as DDBB
from RDU import RepVGGBlock as DRVB
from RDU import DRCB
from RDU import ResDDBB,ResSRB,StaticDCD
import functools

class Unit(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=1, dilation=1, groups=1,
                 internal_channels_1x1_3x3=None, num_experts=4, padding_mode='zeros',
                 deploy=False, dynamic=True, L=None, nonlinear=None, single_init=False, style = 'DBB', res=False, n_channels = None):
        super(Unit,self).__init__()

        #res = true: dynamic = true, style in DBB, VGG
        if n_channels is None:
            n_channels = in_channels

        self.style = style
        self.dynamic = dynamic
        self.res = res

        if style =='DBB':
            self.unit = DDBB(in_channels,out_channels,kernel_size,stride,padding,dilation,groups,internal_channels_1x1_3x3,deploy,dynamic,L,nonlinear,single_init)
        elif style =='VGG':
            self.unit = DRVB(in_channels,out_channels,kernel_size,stride,padding,dilation,groups,padding_mode,nonlinear,deploy,dynamic,L)
        elif style =='CC':
            self.unit = DRCB(in_channels,out_channels,kernel_size,stride,padding,dilation,groups,padding_mode,deploy,num_experts)
        elif style =='ResDBB':
            self.unit = ResDDBB(in_channels,out_channels,n_channels,kernel_size,stride,padding,dilation,groups,internal_channels_1x1_3x3,deploy,dynamic,L,nonlinear,single_init)
        elif style =='SRD':
            self.unit = ResSRB(in_channels,out_channels,kernel_size,stride,padding,dilation,groups,internal_channels_1x1_3x3,deploy,dynamic,L,nonlinear,single_init)
        elif style =='StaticDCD':
            self.unit = StaticDCD(in_channels, out_channels, kernel_size,stride,padding,dilation,groups,dynamic,L,nonlinear)
        else:
            self.unit = nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding,dilation,groups,padding_mode)

        if nonlinear == None:
            self.act = nn.Identity()
        else:
            self.act = nonlinear

    def forward(self, inputs):
        if self.style == 'SRB':
            return self.act(inputs+self.unit(inputs))
        elif self.style == 'Static':
            return self.act(self.unit(inputs))

        if self.dynamic and self.style in {'DBB','VGG','ResDBB','SRD','StaticDCD'}:
            out, res = self.unit(inputs)
            if self.res: return out, res
            else: return out
        else: return self.unit(inputs)

class SeparableConv(nn.Module):
    def __init__(self, n_feats=50, k=3,deploy=True,dynamic=True,L=None,style='DBB',res=False, same = False):
        super(SeparableConv, self).__init__()
        self.res = res

        if not same:
            nf_feats = 2*n_feats
        else:
            nf_feats = n_feats

        if self.res:
             self.conv2 =nn.Conv2d(nf_feats, n_feats, 1, 1, 0)

        self.unit = Unit(n_feats,nf_feats,k,1,(k-1)//2,groups=n_feats,deploy=deploy,dynamic=dynamic,L=L,style=style,res =res)
        self.conv1 =nn.Conv2d(nf_feats, n_feats, 1, 1, 0)

        self.act = nn.LeakyReLU(0.05, True)

    def forward(self, x):
        if self.res:
            out,res = self.unit(x)
            out = self.conv1(self.act(out))
            res = self.conv2(res)
            out += x
            out = self.act(out)
            res = self.act(res)
            return out, res
        else:
            out = self.unit(x)
            out = self.conv1(self.act(out))
            out += x
            out = self.act(out)
            return out

class DFDB(nn.Module):
    def __init__(self, in_channels,deploy=True,dynamic=True,L=None,style='DBB',res=True):
        super(DFDB, self).__init__()
        self.dc = self.distilled_channels = in_channels//2
        self.rc = self.remaining_channels = in_channels
        self.res = res

        self.act = nn.LeakyReLU(0.05, True)

        self.c1_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
        self.c1_r = Unit(self.remaining_channels, self.rc, 3, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act)
        self.c2_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
        self.c2_r = Unit(self.remaining_channels, self.rc, 3, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act)
        self.c3_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
        self.c3_r = Unit(self.remaining_channels, self.rc, 3, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act)
        self.c4 = nn.Conv2d(self.remaining_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
        self.c5 = nn.Conv2d(4*self.dc, in_channels, 1, 1, 0, bias=True, dilation = 1, groups=1)

        if self.res:
            self.PAconv = nn.Conv2d(in_channels, in_channels, 1)
            ##add sconv to enhance res
            #self.sconv = SeparableConv(in_channels,5,deploy,dynamic,L)

            self.conv1x1 = nn.Conv2d(4*in_channels, in_channels, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.sigmoid = nn.Sigmoid()

        self.esa = ESA(in_channels, nn.Conv2d)

    def forward(self, input):
        if self.res:
            distilled_c1 = self.act(self.c1_d(input))
            r_c1,res1 = self.c1_r(input)

            distilled_c2 = self.act(self.c2_d(r_c1))
            r_c2,res2 = self.c2_r(r_c1)

            distilled_c3 = self.act(self.c3_d(r_c2))
            r_c3,res3 = self.c3_r(r_c2)

            r_c4 = self.c4(r_c3)

            out = torch.cat((distilled_c1, distilled_c2, distilled_c3, r_c4), dim=1)

            out_fused = self.c5(out)

            # dfdb res
            scale = self.sigmoid(self.PAconv(out_fused))
            res = self.conv1x1(torch.cat((input, res1, res2, res3), dim=1))
            res = torch.mul(scale, res)

            out_fused = self.esa(out_fused + res)
            #print(out_fused.size())
            return out_fused , res

        else:
            distilled_c1 = self.act(self.c1_d(input))
            r_c1 = self.c1_r(input)

            distilled_c2 = self.act(self.c2_d(r_c1))
            r_c2 = self.c2_r(r_c1)

            distilled_c3 = self.act(self.c3_d(r_c2))
            r_c3 = self.c3_r(r_c2)

            r_c4 = self.c4(r_c3)

            out = torch.cat((distilled_c1, distilled_c2, distilled_c3, r_c4), dim=1)
            out_fused = self.esa(self.c5(out))

            return out_fused

class DFDB_lite(nn.Module):
    def __init__(self, in_channels,deploy=True,dynamic=True,L=None,style='DBB',res=True):
        super(DFDB_lite, self).__init__()
        self.dc = self.distilled_channels = in_channels//2
        self.rc = self.remaining_channels = in_channels
        self.res = res

        if style == 'ResDBB':
            self.act = nn.LeakyReLU(0.05, True)
            self.c1_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c1_r = Unit(self.remaining_channels, self.rc, 7, groups=self.rc, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act,n_channels=2*self.rc)
            self.c2_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c2_r = Unit(self.remaining_channels, self.rc, 5, groups=self.rc, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act,n_channels=2*self.rc)
            self.c3_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c3_r = Unit(self.remaining_channels, self.rc, 3, groups=1, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act,n_channels=2*self.rc)
            self.c4 = nn.Conv2d(self.remaining_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c5 = nn.Conv2d(4*self.dc, in_channels, 1, 1, 0, bias=True, dilation = 1, groups=1)

        else:
            self.act = nn.LeakyReLU(0.05, True)
            self.c1_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c1_r = SeparableConv(self.remaining_channels, 7, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res)
            self.c2_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c2_r = SeparableConv(self.remaining_channels, 5, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res)
            self.c3_d = nn.Conv2d(in_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c3_r = Unit(self.remaining_channels, self.rc, 3, deploy=deploy,dynamic=dynamic,L=L,style=style,res=res,nonlinear=self.act)
            self.c4 = nn.Conv2d(self.remaining_channels, self.dc, 1, 1, 0, bias=True, dilation = 1, groups=1)
            self.c5 = nn.Conv2d(4*self.dc, in_channels, 1, 1, 0, bias=True, dilation = 1, groups=1)

        if self.res:
            self.PAconv = nn.Conv2d(in_channels, in_channels, 1)
            self.conv1x1 = nn.Conv2d(4*in_channels, in_channels, 1, 1, 0, bias=True, dilation = 1, groups=1)
            #self.sconv = SeparableConv(in_channels,5,deploy,dynamic,L,same = True)
            self.sigmoid = nn.Sigmoid()

        self.esa = ESA(in_channels, nn.Conv2d)

    def forward(self, input):
        distilled_c1 = self.act(self.c1_d(input))
        r_c1,res1 = self.c1_r(input)

        distilled_c2 = self.act(self.c2_d(r_c1))
        r_c2,res2 = self.c2_r(r_c1)

        distilled_c3 = self.act(self.c3_d(r_c2))
        r_c3,res3 = self.c3_r(r_c2)

        r_c4 = self.c4(r_c3)

        out = torch.cat((distilled_c1, distilled_c2, distilled_c3, r_c4), dim=1)
        out_fused = self.c5(out)

        if self.res:
            scale = self.sigmoid(self.PAconv(out_fused))
            res = self.conv1x1(torch.cat((input, res1, res2, res3), dim=1))
            res = torch.mul(scale, res)

            out_fused = self.esa(out_fused + res)
            return out_fused , res

        else: return self.esa(out_fused)

def pixelshuffle_block(in_channels, out_channels, upscale_factor=2, kernel_size=3, stride=1):
    #pa = PA(in_channels)
    conv = nn.Conv2d(in_channels, out_channels* (upscale_factor ** 2), 3, 1, 1, bias=True)
    pixel_shuffle = nn.PixelShuffle(upscale_factor)

    return nn.Sequential(conv, pixel_shuffle)

def pixelshuffle_block_cpc(in_channels, out_channels, upscale_factor=2, kernel_size=3, stride=1):
    #pa = PA(in_channels)
    conv = nn.Conv2d(in_channels, out_channels* (upscale_factor ** 2), 3, 1, 1, bias=True)
    pixel_shuffle = nn.PixelShuffle(upscale_factor)
    #new add
    conv1 = nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=True)

    return nn.Sequential(conv, pixel_shuffle, conv1)

class DDistill_SR(nn.Module):
    def __init__(self, in_nc=3, out_nc=3, nf=50, num_modules=4, scale=4,lite=False, deploy=True,dynamic=True,L=None,style='DBB',res=True):
        super(DDistill_SR, self).__init__()

        self.scale = scale
        self.fea_conv = nn.Conv2d(3, nf, 3, 1, 1, bias=True)
        self.res = res

        if lite:
            self.B1 = DFDB_lite(nf,deploy,dynamic,L,style,res)
            self.B2 = DFDB_lite(nf,deploy,dynamic,L,style,res)
            self.B3 = DFDB_lite(nf,deploy,dynamic,L,style,res)
            self.B4 = DFDB_lite(nf,deploy,dynamic,L,style,res)
        else:
            self.B1 = DFDB(nf,deploy,dynamic,L,style,res)
            self.B2 = DFDB(nf,deploy,dynamic,L,style,res)
            self.B3 = DFDB(nf,deploy,dynamic,L,style,res)
            self.B4 = DFDB(nf,deploy,dynamic,L,style,res)

        if self.res:
            self.PAconv = nn.Conv2d(nf, nf, 1)
            self.sigmoid = nn.Sigmoid()
            self.cres = nn.Sequential(
                nn.Conv2d(num_modules*nf, nf, 1, 1, 0, bias=True),
                nn.LeakyReLU(negative_slope=0.1, inplace=True),
            )

        self.c =  nn.Sequential(
            nn.Conv2d(num_modules*nf, nf, 1, 1, 0, bias=True),
            nn.LeakyReLU(negative_slope=0.1, inplace=True),
        )

        self.LR_conv = Unit(nf, nf, 3, deploy=deploy,dynamic=dynamic,L=L,style=style)
        #self.LR_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        #self.esa = ESA(nf, nn.Conv2d)

        upsample_block = pixelshuffle_block_cpc
        #self.upsampler = upsample_block(nf, out_nc, self.scale)
        if self.scale == 2:
            self.upsamplerx2 = upsample_block(nf, out_nc, self.scale)
        elif self.scale ==3:
            self.upsamplerx3 = upsample_block(nf, out_nc, self.scale)
        else:
            self.upsamplerx4 = upsample_block(nf, out_nc, self.scale)
        #self.FConv = nn.Conv2d(out_nc, out_nc, 3, 1, 1, bias=True)

    def forward(self, input):
        out_fea = self.fea_conv(input)

        if self.res:
            out_B1, res1 = self.B1(out_fea)
            out_B2, res2 = self.B2(out_B1)
            out_B3, res3 = self.B3(out_B2)
            out_B4, res4 = self.B4(out_B3)

            out_res = self.cres(torch.cat([res1, res2, res3, res4], dim=1))

            out_B = self.c(torch.cat([out_B1, out_B2, out_B3, out_B4], dim=1))
            #out_B = self.esa(out_B)

            scale = self.sigmoid(self.PAconv(out_B))
            out_res = torch.mul(scale, out_res)

            out_lr = self.LR_conv(out_B + out_res) + out_fea

        else:
            out_B1 = self.B1(out_fea)
            out_B2 = self.B2(out_B1)
            out_B3 = self.B3(out_B2)
            out_B4 = self.B4(out_B3)

            out_B = self.c(torch.cat([out_B1, out_B2, out_B3, out_B4], dim=1))
            #out_B = self.esa(out_B)
            out_lr = self.LR_conv(out_B) + out_fea

        if self.scale == 2:
            output = self.upsamplerx2(out_lr)
        elif self.scale == 3:
            output = self.upsamplerx3(out_lr)
        else:
            output = self.upsamplerx4(out_lr)

        return output.clamp(0,1)

def model_convert(model:torch.nn.Module, save_path=None, do_copy=True):
    if do_copy:
        model = copy.deepcopy(model)
    for module in model.modules():
        if hasattr(module, 'switch_to_deploy'):
            module.switch_to_deploy()
    if save_path is not None:
        torch.save(model.state_dict(), save_path)
    return model

if __name__ == "__main__":

    model = DDistill_SR(in_nc=3, out_nc=3, nf = 56, num_modules=4, scale=4, deploy=False, dynamic=True, L = 16, style = 'DBB')
    model.load_state_dict(torch.load(r'model-backup/DDistillSRx4.pth'))

    from PIL import Image
    import torchvision.transforms as transforms

    model.eval()

    model_inference = model_convert(model)

    img = Image.open(r'Set5/LR_bicubic/X4/babyx4.png').convert('RGB')

    x = transforms.ToTensor()(img)
    x = x.unsqueeze(0)
    out1 = model(x)
    out2 = model_inference(x)
    #print(model)

    img1 = transforms.ToPILImage()(out1.squeeze(0))
    img1.show()
    img2 = transforms.ToPILImage()(out2.squeeze(0))
    img2.show()
    f = torch.nn.L1Loss()

    z = f(out1,out2)
    print(z)


#!/usr/bin/env python

"""
Computes the budget of Pb, Be7, and Be10 from 1-year
TransportTracersBenchmark simulations.
"""

# ======================================================================
# Imports etc.
# ======================================================================

import os
from glob import glob
import warnings
from calendar import monthrange
import numpy as np
import xarray as xr
import gcpy.constants as constants
from gcpy.grid import get_troposphere_mask
import gcpy.util as util
import gc

# Suppress harmless run-time warnings (mostly about underflow in division)
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=UserWarning)

# ======================================================================
# Define a class for passing global variables to the methods below
# ======================================================================

class _GlobVars:
    """
    Private class _GlobVars contains global data that needs to be
    shared among the methods in this module.
    """

    def __init__(self, devstr, devdir, devrstdir, year, dst, is_gchp,
                 gchp_res, gchp_is_pre_14_0, overwrite, spcdb_dir):
        """
        Initializes the _GlobVars class.

        Args:
            devstr: str
                Label denoting the "Dev" version.
            devdir: str
                Directory where diagnostic files are found.
            devrstdir: str
                Directory where restart files are found.
            dst: str
                Directory where plots & tables will be created.
            year: int
                Year of the benchmark simulation.
            is_gchp: bool
                Denotes if this is GCHP (True) or GCC (False) data.
            gchp_res: str
                A string (e.g. "c24") denoting GCHP grid resolution.
            gchp_is_pre_14_0: bool
                Logical to indicate whether or not the GCHP data is prior
                to GCHP 14.0.0.  Needed for restart files only.
            overwrite: bool
                Denotes whether to ovewrite existing budget tables.
            spcdb_dir: str
                Directory where species_database.yml is stored.
        """
        # ------------------------------
        # Arguments from outside
        # ------------------------------
        self.devstr = devstr
        self.devdir = devdir
        self.devrstdir = devrstdir
        self.dst = dst
        self.is_gchp = is_gchp
        self.gchp_res = gchp_res
        self.gchp_is_pre_14_0 = gchp_is_pre_14_0
        self.overwrite = overwrite

        # ------------------------------
        # Benchmark year
        # ------------------------------
        self.y0 = year
        self.y1 = self.y0 + 1
        self.y0_str = "{}".format(self.y0)
        self.y1_str = "{}".format(self.y1)

        # ------------------------------
        # Collection file lists
        # ------------------------------

        # Initial restart file
        RstInit = util.get_filepath(
            self.devrstdir,
            "Restart",
            np.datetime64(f"{self.y0_str}-01-01T00:00:00"),
            is_gchp=self.is_gchp,
            gchp_res=self.gchp_res,
            gchp_is_pre_14_0=self.gchp_is_pre_14_0
        )

        # Final restart file
        RstFinal = util.get_filepath(
            self.devrstdir,
            "Restart",
            np.datetime64(f"{self.y1_str}-01-01T00:00:00"),
            is_gchp=self.is_gchp,
            gchp_res=self.gchp_res,
            gchp_is_pre_14_0=self.gchp_is_pre_14_0
        )

        # Diagnostics
        HemcoDiag = os.path.join(
            self.devdir,
            f"HEMCO_diagnostics.{self.y0_str}*.nc"
        )

        DryDep = os.path.join(
            self.devdir,
            f"*.DryDep.{self.y0_str}*.nc4"
        )

        RadioNucl = os.path.join(
            self.devdir,
            f"*.RadioNuclide.{self.y0_str}*.nc4"
        )

        if is_gchp:
            StateMetAvg = os.path.join(
                self.devdir,
                f"*.StateMet_avg.{self.y0_str}*.nc4"
            )

            StateMet = os.path.join(
                self.devdir,
                f"*.StateMet.{self.y0_str}*.nc4"
            )

            # Set a logical if we need to read StateMet_avg or StateMet
            gchp_use_statemet_avg  = os.path.exists(StateMetAvg)

        else:
            StateMet = os.path.join(
                self.devdir,
                f"*.StateMet.{self.y0_str}*.nc4"
            )

        SpeciesConc = os.path.join(
            self.devdir,
            f"*.SpeciesConc.{self.y0_str}*.nc4"
        )

        WetLossConv = os.path.join(
            self.devdir,
            f"*.WetLossConv.{self.y0_str}*.nc4"
        )

        WetLossLS = os.path.join(
            self.devdir,
            f"*.WetLossLS.{self.y0_str}*.nc4"
        )

        GCHPEmiss = os.path.join(
            self.devdir,
            f"*.Emissions.{self.y0_str}*.nc4"
        )

        # ------------------------------
        # Read data collections
        # ------------------------------

        # Restarts
        extra_kwargs = {}

        self.ds_ini = xr.open_dataset(
            RstInit,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        self.ds_end = xr.open_dataset(
            RstFinal,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        # Change the restart datasets into format similar to GCC, and flip
        # vertical axis.  Also test if the restart files have the BXHEIGHT
        # variable contained within them.
        if is_gchp:
            self.ds_ini = util.rename_and_flip_gchp_rst_vars(self.ds_ini)
            self.ds_end = util.rename_and_flip_gchp_rst_vars(self.ds_end)

        # Diagnostics
        self.ds_dcy = xr.open_mfdataset(
            RadioNucl,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        self.ds_dry = xr.open_mfdataset(
            DryDep,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        self.ds_cnc = xr.open_mfdataset(
            SpeciesConc,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        self.ds_wcv = xr.open_mfdataset(
            WetLossConv,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        self.ds_wls = xr.open_mfdataset(
            WetLossLS,
            drop_variables=constants.skip_these_vars,
            **extra_kwargs
        )

        # Met fields
        # For GCHP: Read from StateMet_avg if present (otherwise StateMet)
        if is_gchp and gchp_use_statemet_avg:
            self.ds_met = xr.open_mfdataset(
                StateMetAvg,
                drop_variables=constants.skip_these_vars,
                **extra_kwargs
            )
        else:
            self.ds_met = xr.open_mfdataset(
                StateMet,
                drop_variables=constants.skip_these_vars,
                **extra_kwargs
            )

        # Emissions
        if self.is_gchp:
            self.ds_hco = xr.open_mfdataset(
                GCHPEmiss,
                drop_variables=constants.skip_these_vars,
                **extra_kwargs
            )
        else:
            self.ds_hco = xr.open_mfdataset(
                HemcoDiag,
                drop_variables=constants.skip_these_vars,
                **extra_kwargs
            )

        # Area and troposphere mask
        # The area in m2 is on the restart file grid
        # The area in cm2 is on the History diagnostic grid
        # Both grids are identical in GCClassic but differ in GCHP
        if self.is_gchp:
            if 'Met_AREAM2' not in self.ds_met.data_vars.keys():
                msg = 'Could not find Met_AREAM2 in StateMet_avg collection!'
                raise ValueError(msg)
            area_m2 = self.ds_met["Met_AREAM2"].isel(time=0)
            area_m2 = util.reshape_MAPL_CS(area_m2)
            self.area_m2 = area_m2
            self.area_cm2 = self.ds_met["Met_AREAM2"] * 1.0e4
        else:
            self.area_m2 = self.ds_met["AREA"].isel(time=0)
            self.area_cm2 = self.area_m2 * 1.0e4
        self.tropmask = get_troposphere_mask(self.ds_met)

        # ------------------------------
        # Months and days
        # ------------------------------
        self.N_MONTHS = 12
        self.N_MONTHS_FLOAT = self.N_MONTHS * 1.0

        # Days per month in the benchmark year
        self.d_per_mon = np.zeros(self.N_MONTHS)
        for t in range(self.N_MONTHS):
            self.d_per_mon[t] = monthrange(self.y0, t + 1)[1] * 1.0

        # Days in the benchmark year
        self.d_per_yr = np.sum(self.d_per_mon)

        # Fraction of year occupied by each month
        self.frac_of_yr = np.zeros(self.N_MONTHS)
        for t in range(self.N_MONTHS):
            self.frac_of_yr[t] = self.d_per_mon[t] / self.d_per_yr

        # ------------------------------
        # Species info
        # ------------------------------

        # List of species (and subsets for the trop & strat)
        self.species_list = ["Pb210", "Be7", "Be10"]

        # Read the species database
        path = os.path.join(spcdb_dir, "species_database.yml")
        spcdb = util.read_config_file(path)

        # Molecular weights [g mol-1], as taken from the species database
        self.mw = {}
        for v in self.species_list:
            self.mw[v] = spcdb[v]["MW_g"]
        self.mw["Air"] = constants.MW_AIR_g

        # kg/s --> g/day
        self.kg_s_to_g_d_value = 86400.0 * 1000.0

        # ------------------------------
        # Conversion factors
        # ------------------------------
        self.kg_per_mol = {}
        self.vv_to_g = {}
        self.mcm2s_to_g_d = {}
        self.kg_s_to_g_d = {}
        for spc in self.species_list:

            # kg/s --> g/day (same for all species)
            self.kg_s_to_g_d[spc] = self.kg_s_to_g_d_value

            # kg/mole for each species
            self.kg_per_mol[spc] = constants.AVOGADRO / (self.mw[spc] * 1e-3)

            # v/v dry --> g
            self.vv_to_g[spc] = self.ds_met["Met_AD"].values  \
                * (self.mw[spc] / self.mw["Air"]) * 1000.0

            # molec/cm2/s --> g/day
            self.mcm2s_to_g_d[spc] = self.area_cm2.values \
                / self.kg_per_mol[spc] \
                * self.kg_s_to_g_d[spc]

# ======================================================================
# Functions
# ======================================================================

def total(globvars, dict_list):
    """
    Function to take the difference of two dict objects.
    Assumes that all objects have the same keys.

    Args:
        globvars: obj of type _GlobVars
            Global variables needed for budget computations.
        dict_list: list of dict
            Dictionaries to be summed.

    Returns:
        result: dict
            Key-by-key sum of all dicts in dict_list.
    """
    # Initialize
    result = {}
    for spc in globvars.species_list:
        result[spc + "_f"] = 0.0   # full-atmosphere
        result[spc + "_t"] = 0.0   # trop-only
        result[spc + "_s"] = 0.0   # strat-only

    # Sum over all dictionaries
    for d in dict_list:
        for k, v in d.items():
            result[k] += v

    return result

def mass_from_rst(globvars, ds, tropmask):
    """
    Computes global species mass from a restart file.

    Args:
        globvars: obj of type _GlobVars
            Global variables needed for budget computations.
        ds: xarray Dataset
            Data containing species mass to be summed.
        tropmask: numpy ndarray
            Mask to denote tropospheric grid boxes.

    Returns:
        result: dict
            Species mass in strat, trop, and strat+trop regimes.
    """
    # Initialize
    vv_to_g = {}
    rst_f = {}
    rst_t = {}
    result = {}

    # Conversion factors based on restart-file met fields
    # NOTE: Convert DataArray to numpy ndarray so that the
    # broadcasting will be done properly!!!
    g100 = 100.0 / constants.G
    if 'time' in ds["Met_DELPDRY"].dims:
        deltap = ds["Met_DELPDRY"].isel(time=0).values
    else:
        deltap = ds["Met_DELPDRY"].values
    area = globvars.area_m2.values
    airmass = deltap * area * g100

    # GCClassic: Restart variables begin with SpeciesRst
    prefix = "SpeciesRst_"

    # Determine if restart file variables have a time dimension
    # (if one does, they all do)
    varname = prefix + globvars.species_list[0]
    if 'time' in ds[varname].dims:
        has_time_dim = True
    else:
        has_time_dim = False

    # Loop over species
    for spc in globvars.species_list:

        # Add the prefix to the species name
        varname = prefix + spc

        # Conversion factor from mixing ratio to g, w/ met from rst file
        vv_to_g[spc] = airmass \
            * (globvars.mw[spc] / globvars.mw["Air"]) * 1000.0

        if has_time_dim:
            rst_f[spc] = ds[varname].isel(time=0).values * vv_to_g[spc]
        else:
            rst_f[spc] = ds[varname].values * vv_to_g[spc]

        # Troposphere-only mass
        rst_t[spc] = np.ma.masked_array(rst_f[spc], tropmask)

        # Sums
        result[spc + "_f"] = np.sum(rst_f[spc])
        result[spc + "_t"] = np.sum(rst_t[spc])
        result[spc + "_s"] = result[spc + "_f"] - result[spc + "_t"]

    return result

def annual_average(globvars, ds, collection, conv_factor):
    """
    Computes the annual average of budgets or fluxes.

    Args:
        globvars: obj of type _GlobVars
            Global variables needed for budget computations.
        ds: xarray Dataset
            Data to be averaged
        collection: str
            Name of the diagnostic collection.
        conv_factor: str
            Conversion factor to be applied.

     Returns:
        result: dict
            Annual-average budgets or fluxes in
            in strat, trop, and strat+trop regimes.
    """

    # Initialize
    q = {}
    q_sum_f = np.zeros(globvars.N_MONTHS)
    q_sum_t = np.zeros(globvars.N_MONTHS)
    q_sum_s = np.zeros(globvars.N_MONTHS)
    result = {}

    for spc in globvars.species_list:

        # Whole-atmosphere quanity [g] or [g d-1]
        if "SpeciesConc" in collection:
            varname = collection.strip() + "VV_" + spc
        else:
            varname = collection.strip() + "_" + spc
        q[spc + "_f"] = ds[varname].values * conv_factor[spc]

        # Shape of the data
        #q_shape = q[spc + "_f"].shape

        if "DryDep" in collection:
            # NOTE: DryDep is by nature trop-only
            # Therefore, data arrays don't have a lev dimension,
            # so special handling must be done.
            if globvars.is_gchp:
                sum_axes = (1, 2, 3)
            else:
                sum_axes = (1, 2)

        else:
            # Otherwise, expect a lev dimension
            if globvars.is_gchp:
                sum_axes = (1, 2, 3, 4)
            else:
                sum_axes = (1, 2, 3)

            # Trop-only quantities [g] or [g d-1]
            q[spc + "_t"] = np.ma.masked_array(q[spc + "_f"],
                                               globvars.tropmask)

        # Compute monthly averages, weighted by # of days in month
        # Special handling for DryDep, which lacks a "lev" dimension.
        if "DryDep" in collection:
            q_sum_f = np.sum(q[spc + "_f"], axis=sum_axes) \
                * globvars.d_per_mon
            q_sum_t = q_sum_f
            q_sum_s = 0.0
        else:
            q_sum_f = np.sum(q[spc + "_f"], axis=sum_axes) \
                * globvars.d_per_mon
            q_sum_t = np.sum(q[spc + "_t"], axis=sum_axes) \
                * globvars.d_per_mon
            q_sum_s = q_sum_f - q_sum_t

        # Take annual averages
        result[spc + "_f"] = np.sum(q_sum_f) / globvars.d_per_yr
        result[spc + "_t"] = np.sum(q_sum_t) / globvars.d_per_yr
        result[spc + "_s"] = np.sum(q_sum_s) / globvars.d_per_yr

    return result

def annual_average_sources(globvars):
    """
    Computes the annual average of radionuclide sources.

    Args:
        globvars: obj of type _GlobVars
            Global variables needed for budget computations.

     Returns:
        result: dict
            Source totals in strat, trop, and strat+trop regimes.
    """

    # Initialize
    q = {}
    q_sum_f = np.zeros(globvars.N_MONTHS)
    q_sum_t = np.zeros(globvars.N_MONTHS)
    q_sum_s = np.zeros(globvars.N_MONTHS)
    result = {}

    # Pb210 source is from Rn222 decay, in the RadioNuclide collection
    q["Pb210_f"] = globvars.ds_dcy["PbFromRnDecay"].values \
        * globvars.kg_s_to_g_d["Pb210"]

    # Determine the shape of the data
    q_shape = q["Pb210_f"].shape     # e.g. (time,lev,lat,lon)
    n_levs = q_shape[1]              # Number of levels

    # Determine which array axes to sum over for monthly sums
    if globvars.is_gchp:
        area_var = "Met_AREAM2"
        sum_axes = (1, 2, 3, 4)
    else:
        area_var = "AREA"
        sum_axes = (1, 2, 3)

    # Be7 and Be10 sources are in the HEMCO diagnostics collection
    q["Be7_f"] = np.zeros(q_shape)
    q["Be10_f"] = np.zeros(q_shape)

    # Error traps
    if q_shape[0] != globvars.N_MONTHS:
        msg = f"Expected {globvars.N_MONTHS} but only found {q_shape[0]}"
        raise ValueException(msg)
    if q_shape[1] != n_levs:
        msg = f"Expected {n_levs} but only found {q_shape[1]}!"
        raise ValueException(msg)

    # Convert Be7 and Be10 sources from kg/m2/s to g/day
    # If GCHP data, must vertically flip the emissions diagnostic
    for t in range(globvars.N_MONTHS):
        for k in range(n_levs):
            if globvars.is_gchp:
                kf = n_levs - k - 1
                q["Be7_f"][t, k, :, :, :] = \
                    globvars.ds_hco["EmisBe7_Cosmic"].isel(time=t, lev=kf) * \
                    globvars.ds_met[area_var].isel(time=t) * \
                    globvars.kg_s_to_g_d["Be7"]
                q["Be10_f"][t, k, :, :, :] = \
                    globvars.ds_hco["EmisBe10_Cosmic"].isel(time=t, lev=kf) * \
                    globvars.ds_met[area_var].isel(time=t) * \
                    globvars.kg_s_to_g_d["Be10"]
            else:
                q["Be7_f"][t, k, :, :] = \
                    globvars.ds_hco["EmisBe7_Cosmic"].isel(time=t, lev=k) * \
                    globvars.ds_met[area_var].isel(time=t) * \
                    globvars.kg_s_to_g_d["Be7"]
                q["Be10_f"][t, k, :, :] = \
                    globvars.ds_hco["EmisBe10_Cosmic"].isel(time=t, lev=k) * \
                    globvars.ds_met[area_var].isel(time=t) * \
                    globvars.kg_s_to_g_d["Be10"]

    # Now proceed to computing the annual averages
    for spc in globvars.species_list:

        # Tropospheric-only quantities
        q[spc + "_t"] = np.ma.masked_array(q[spc + "_f"], globvars.tropmask)

        # Take monthly sums, weighted by the # of days in the month
        q_sum_f = np.sum(q[spc + "_f"], axis=sum_axes) \
            * globvars.d_per_mon
        q_sum_t = np.sum(q[spc + "_t"], axis=sum_axes) \
            * globvars.d_per_mon
        q_sum_s = q_sum_f - q_sum_t

        # Take annual averages
        result[spc + "_f"] = np.sum(q_sum_f) / globvars.d_per_yr
        result[spc + "_t"] = np.sum(q_sum_t) / globvars.d_per_yr
        result[spc + "_s"] = np.sum(q_sum_s) / globvars.d_per_yr

    return result

def trop_residence_time(globvars):
    """
    Computes the tropospheric residence time of radionuclides.

    Args:
        globvars: obj of type _GlobVars
            Global variables needed for budget computations.

    Returns:
        result: dict
            Tropopsheric residence time for all species.
    """

    # Initialize
    result = {}

    # Loop over species
    for spc in globvars.species_list:

        # Initialize
        result[spc + "_t"] = 0.0

        # Concentration [g]
        var = "SpeciesConcVV_" + spc
        q_cnc = globvars.ds_cnc[var].values * globvars.vv_to_g[spc]
        q_cnc = np.ma.masked_array(q_cnc, globvars.tropmask)

        # DryDep [g d-1]
        var = "DryDep_" + spc
        q_dry = globvars.ds_dry[var].values * globvars.mcm2s_to_g_d[spc]

        # Convective wet scavenging [g d-1]
        var = "WetLossConv_" + spc
        q_wcv = globvars.ds_wcv[var].values * globvars.kg_s_to_g_d[spc]
        q_wcv = np.ma.masked_array(q_wcv, globvars.tropmask)

        # Large-scale wet scavenging [g d-1]
        var = "WetLossLS_" + spc
        q_wls = globvars.ds_wls[var].values * globvars.kg_s_to_g_d[spc]
        q_wls = np.ma.masked_array(q_wls, globvars.tropmask)

        # Set a flag to denote if this is GCHP
        # and also denote which axes we should sum over
        if globvars.is_gchp:
            sum_axes = (1, 2, 3, 4)
            sum_dryd = (1, 2, 3)
        else:
            sum_axes = (1, 2, 3)
            sum_dryd = (1, 2)

        # Compute monthly averages [g]
        q_cnc_sum = np.sum(q_cnc, axis=sum_axes)
        q_dry_sum = np.sum(q_dry, axis=sum_dryd)
        q_wcv_sum = np.sum(q_wcv, axis=sum_axes)
        q_wls_sum = np.sum(q_wls, axis=sum_axes)

        # Compute lifetime
        for t in range(globvars.N_MONTHS):
            num = q_cnc_sum[t]
            denom = q_dry_sum[t] + q_wcv_sum[t] + q_wls_sum[t]
            result[spc + "_t"] += 1.0 / (num / denom)

        # Convert residence time [d]
        result[spc + "_t"] = 1.0 \
            / (result[spc + "_t"] / globvars.N_MONTHS_FLOAT)

    return result

def print_budgets(globvars, data, key):
    """
    Prints the trop+strat budget file.

    Args:
        globvars: object of type _GlobVars
            Global variables needed for budget computations.
        data: dict
            Nested dictionary containing budget info.
        key: list of str
            One of "_f", (full-atmosphere) "_t" (trop-only),
            or "_s" (strat-only).
    """
    # Create the plot directory hierarchy if it doesn't already exist
    if os.path.isdir(globvars.dst) and not globvars.overwrite:
        err_str = "Pass overwrite=True to overwrite files in that directory"
        print(f"Directory {globvars.dst} exists. {err_str}")
        return
    elif not os.path.isdir(globvars.dst):
        os.makedirs(globvars.dst)

    # Filename to print
    if "_f" in key:
        filename = f"{globvars.dst}/Pb-Be_budget_trop_strat.txt"
    elif "_t" in key:
        filename = f"{globvars.dst}/Pb-Be_budget_troposphere.txt"
    elif "_s" in key:
        filename = f"{globvars.dst}/Pb-Be_budget_stratosphere.txt"

    # Common title string
    title = "Annual Average Global Budgets of 210Pb, 7Be, and 10Be\n        "

    # Open file and print budgets
    with open(filename, "w+") as f:
        if "_f" in key:
            print(
                f"Table 1. {title} in the Troposphere + Stratosphere in {globvars.devstr} for year {globvars.y0_str}\n", file=f)
        elif "_t" in key:
            print(f"Table 2. {title} in the Troposphere in {globvars.devstr} for year {globvars.y0_str}\n", file=f)
        elif "_s" in key:
            print(f"Table 3. {title} in the Stratosphere in {globvars.devstr} for year {globvars.y0_str}\n", file=f)
        print(
            "                                210Pb          7Be         10Be",
            file=f)
        print(
            "                          -----------  -----------  -----------",
            file=f)

        vals = [v[1] for v in list(data["burden"].items()) if key in v[0]]
        print("  Burden, g               {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)
        print(file=f)

        if "_t" in key:
            vals = [v[1]
                    for v in list(data["res_time"].items()) if key in v[0]]

            print(
                "  Residence time, d       {:11.4f}  {:11.4f}  {:11.4f}".format(
                    *vals), file=f)
            print(file=f)

        vals = [v[1] for v in list(data["src_total"].items()) if key in v[0]]
        print("  Sources, g d-1          {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)
        print(file=f)

        vals = [v[1] for v in list(data["snk_total"].items()) if key in v[0]]
        print("  Sinks, g d-1            {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)

        vals = [v[1] for v in list(data["snk_drydep"].items()) if key in v[0]]
        print("    Dry deposition        {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)
        print("    Wet deposition", file=f)

        vals = [v[1] for v in list(data["snk_wetls"].items()) if key in v[0]]
        print("      Stratiform          {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)

        vals = [v[1] for v in list(data["snk_wetconv"].items()) if key in v[0]]
        print("      Convective          {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)

        vals = [v[1] for v in list(data["snk_decay"].items()) if key in v[0]]
        print("    Radioactive decay     {:11.7f}  {:11.7f}  {:11.8f}".format(
            *vals), file=f)
        print(file=f)

        vals = [v[1] for v in list(data["src_minus_snk"].items())
                if key in v[0]]
        print("  Sources - Sinks, g d-1  {:11.6f}  {:11.6f}  {:11.6f}".format(
            *vals), file=f)

        print(file=f)
        print("  Accumulation Term", file=f)

        vals = [v[1] for v in list(data["accum_init"].items()) if key in v[0]]
        print("    Initial mass, g       {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)

        vals = [v[1] for v in list(data["accum_final"].items()) if key in v[0]]
        print("    Final mass, g         {:11.4f}  {:11.4f}  {:11.4f}".format(
            *vals), file=f)

        vals = [v[1] for v in list(data["accum_diff"].items()) if key in v[0]]
        print("    Difference, g         {:11.7f}  {:11.7f}  {:11.7f}".format(
            *vals), file=f)
        f.close()

def transport_tracers_budgets(
        devstr,
        devdir,
        devrstdir,
        year,
        dst='./1yr_benchmark',
        is_gchp=False,
        gchp_res="c00",
        gchp_is_pre_14_0=False,
        overwrite=True,
        spcdb_dir=os.path.dirname(__file__)):
    """
    Main program to compute TransportTracersBenchmark budgets

    Args:
        maindir: str
            Top-level benchmark folder
        devstr: str
            Denotes the "Dev" benchmark version.
        year: int
            The year of the benchmark simulation (e.g. 2016).

    Keyword Args (optional):
        dst: str
            Directory where budget tables will be created.
            Default value: './1yr_benchmark'
        is_gchp: bool
            Denotes if data is from GCHP (True) or GCC (false).
            Default value: False
        gchp_res: str
            A string (e.g. "c24") denoting GCHP grid resolution.
            Default value: "c00".
        gchp_is_pre_14_0: bool
            Logical to indicate whether or not the GCHP data is prior
            to GCHP 14.0.0.  Needed for restart files only.
            Default value: False
        overwrite: bool
            Denotes whether to ovewrite existing budget tables.
            Default value: True
        spcdb_dir: str
            Directory where species_database.yml is stored.
            Default value: GCPy directory
    """

    # Store global variables in a private class
    globvars = _GlobVars(
        devstr,
        devdir,
        devrstdir,
        year,
        dst,
        is_gchp,
        gchp_res,
        gchp_is_pre_14_0,
        overwrite,
        spcdb_dir
    )

    # Data structure for budgets
    data = {}

    # ==================================================================
    # Get the accumulation term (init & final masses) from the restarts
    # ==================================================================

    # Get initial mass from restart file
    ds = globvars.ds_ini
    tropmask = globvars.tropmask[0, :, :, :]
    data["accum_init"] = mass_from_rst(
        globvars,
        ds,
        tropmask
    )

    # Get initial mass from restart file
    ds = globvars.ds_end
    tropmask = globvars.tropmask[globvars.N_MONTHS - 1, :, :, :]
    data["accum_final"] = mass_from_rst(
        globvars,
        ds, tropmask
    )

    # Take the difference final - init
    data["accum_diff"] = util.dict_diff(
        data["accum_init"],
        data["accum_final"]
    )

    # ==================================================================
    # Burdens [g]
    # ==================================================================
    data["burden"] = annual_average(
        globvars,
        ds=globvars.ds_cnc,
        collection='SpeciesConc',
        conv_factor=globvars.vv_to_g
    )

    # ==================================================================
    # Sources and sinks [g d-1]
    # ==================================================================
    data["src_total"] = annual_average_sources(globvars)

    # Radioactive decay [g d-1]
    data["snk_decay"] = annual_average(
        globvars,
        ds=globvars.ds_dcy,
        collection="RadDecay",
        conv_factor=globvars.kg_s_to_g_d
    )

    # Drydep fluxes
    data["snk_drydep"] = annual_average(
        globvars,
        ds=globvars.ds_dry,
        collection="DryDep",
        conv_factor=globvars.mcm2s_to_g_d
    )

    # Convective wet scavenging
    data["snk_wetconv"] = annual_average(
        globvars,
        ds=globvars.ds_wcv,
        collection="WetLossConv",
        conv_factor=globvars.kg_s_to_g_d
    )

    # Large-scale wet scavenging [g d-1]
    data["snk_wetls"] = annual_average(
        globvars,
        ds=globvars.ds_wls,
        collection="WetLossLS",
        conv_factor=globvars.kg_s_to_g_d
    )

    # Total sinks
    data["snk_total"] = total(
        globvars,
        [data["snk_drydep"],
         data["snk_wetls"],
         data["snk_wetconv"],
         data["snk_decay"]]
    )

    # Sources - sinks
    data["src_minus_snk"] = util.dict_diff(
        data["snk_total"],
        data["src_total"]
    )

    # Tropospheric residence time
    data["res_time"] = trop_residence_time(globvars)

    # ==================================================================
    # Print budgets
    # ==================================================================
    for key in ["_f", "_t", "_s"]:
        print_budgets(
            globvars,
            data,
            key
        )

    # ==================================================================
    # Force garbage collection
    # ==================================================================
    del globvars
    del data
    del ds
    gc.collect()