"""
Procs for plink 1.9
"""
from pyppl import Proc
from diot import Diot
from . import params, proc_factory
from .vcf import pVcf2Plink
from .tcgamaf import pGTMat2Plink

pPlinkFromVcf = pVcf2Plink.copy()
pPlinkFromGTMat = pGTMat2Plink.copy()

pPlinkStats = proc_factory(
	desc = 'Do basic statistics with plink 1.9',
	config = Diot(annotate = """
	@name:
		pPlinkStats
	"""))
pPlinkStats.input         = 'indir:dir'
pPlinkStats.output        = 'outdir:dir:{{i.indir | fn}}.plinkStats'
pPlinkStats.args.plink    = params.plink.value
pPlinkStats.args.nthread  = False
pPlinkStats.args.params   = Diot({
	'hardy'    : True,
	'het'      : True,
	'freq'     : True,
	'missing'  : True,
	'check-sex': True,
})
pPlinkStats.args.cutoff   = Diot({
	'hardy.hwe'     : 1e-5,
	'hardy.mingt'   : None,
	'het'           : 3,
	'freq'          : 0.01,
	'missing.sample': .95,
	'missing.snp'   : .95,
})
pPlinkStats.args.plot     = Diot({
	'hardy.hwe'     : True,
	'hardy.mingt'   : True,
	'het'           : True,
	'freq'          : True,
	'missing.sample': True,
	'missing.snp'   : True,
})
pPlinkStats.args.devpars  = Diot(res=300, width=2000, height=2000)
pPlinkStats.envs.rimport  = rimport
pPlinkStats.lang          = params.Rscript.value

pPlinkSampleFilter = proc_factory(
	desc = 'Do sample filtering or extraction using `--keep[-fam]` or `--remove[-fam]`',
	config = Diot(annotate = """
	@name:
		pPlinkSampleFilter
	"""))
pPlinkSampleFilter.input        = 'indir:dir, samfile:file'
pPlinkSampleFilter.output       = 'outdir:dir:{{i.indir | bn}}'
pPlinkSampleFilter.args.plink   = params.plink.value
pPlinkSampleFilter.args.keep    = True
pPlinkSampleFilter.args.samid   = 'iid' # both or fid
pPlinkSampleFilter.args.fam     = False
pPlinkSampleFilter.args.params  = Diot()
pPlinkSampleFilter.args.nthread = False
pPlinkSampleFilter.lang         = params.python.value

pPlinkMiss = proc_factory(
	desc = 'Find samples and snps with missing calls',
	config = Diot(annotate = """
	@name:
		pPlinkMiss
	@description:
		Find samples and snps with missing calls, calculate the call rates and plot them.
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outdir:dir`: The output directory. Default: `{{i.indir | fn}}.miss`
			- `.imiss`: The miss calls for samples
			- `.lmiss`: The miss calls for snps
			- `.samplecr.fail`: The samples fail sample call rate cutoff (`args.samplecr`)
			- `.snpcr.fail`: The SNPs fail snp call rate cutoff (`args.snpcr`)
	@args:
		`plink`: The path to plink.
		`samplecr`: The sample call rate cutoff. Default: `.95`
		`snpcr`: The SNP call rate cutoff. Default: `.95`
		`plot`: Whether plot the distribution of the call rates? Default: `True`
		`devpars`: The device parameters for the plot. Default: `Diot(res=300, width=2000, height=2000)`
	"""))
pPlinkMiss.input         = 'indir:dir'
pPlinkMiss.output        = 'outdir:dir:{{i.indir | fn}}.miss'
pPlinkMiss.args.plink    = params.plink.value
pPlinkMiss.args.samplecr = .95
pPlinkMiss.args.snpcr    = .95
pPlinkMiss.args.plot     = True
pPlinkMiss.args.devpars  = Diot(res=300, width=2000, height=2000)
pPlinkMiss.envs.rimport  = rimport
pPlinkMiss.lang          = params.Rscript.value

pPlinkFreq = proc_factory(
	desc = 'Filter snps with minor allele frequency.',
	config = Diot(annotate = """
	@name:
		pPlinkFreq
	"""))
pPlinkFreq.input        = 'indir:dir'
pPlinkFreq.output       = 'outdir:dir:{{i.indir | fn}}.freq'
pPlinkFreq.args.plink   = params.plink.value
pPlinkFreq.args.cutoff  = 0.01
pPlinkFreq.args.plot    = True
pPlinkFreq.args.devpars = Diot(res=300, width=2000, height=2000)
pPlinkFreq.lang         = params.python.value

pPlinkSexcheck = proc_factory(
	desc = 'Check inconsistency between sex denoted and from genotypes.',
	config = Diot(annotate = """
	@name:
		pPlinkSexcheck
	@description:
		Check inconsistency between sex denoted and from genotypes.
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outdir:dir`: The output directory. Default: `{{i.indir | fn}}.sexcheck`
			- `.sexcheck`: The orginal sex check report from `plink`
			- `.sex.fail`: The samples that fail sex check.
	@args:
		`plink`: The path to plink.
	"""))
pPlinkSexcheck.input        = 'indir:dir'
pPlinkSexcheck.output       = 'outdir:dir:{{i.indir | fn}}.sexcheck'
pPlinkSexcheck.args.plink   = params.plink.value
pPlinkSexcheck.envs.rimport = rimport
pPlinkSexcheck.lang         = params.Rscript.value

pPlinkHet = proc_factory(
	desc = 'Calculate the heterozygosity of each sample',
	config = Diot(annotate = """
	@name:
		pPlinkHet
	@description:
		Calculate the heterozygosity of each sample.
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outdir:dir`: The output directory. Default: `{{i.indir | fn}}.het`
			- `.het`: The heterozygosity file generated by `plink`.
			- `.het.fail`: The samples fail sample heterozygosity cutoff (`args.cutoff`)
	@args:
		`plink`: The path to plink.
		`cutoff`: The sample heterozygosity cutoff. Default: `3` (mean-3SD ~ mean+3SD)
		`plot`: Whether plot the distribution of the heterozygosity? Default: `True`
		`devpars`: The device parameters for the plot. Default: `Diot(res=300, width=2000, height=2000)`
	"""))
pPlinkHet.input        = 'indir:dir'
pPlinkHet.output       = 'outdir:dir:{{i.indir | fn}}.het'
pPlinkHet.args.plink   = params.plink.value
pPlinkHet.args.cutoff  = 3
pPlinkHet.args.plot    = True
pPlinkHet.args.devpars = Diot(res=300, width=2000, height=2000)
pPlinkHet.envs.rimport = rimport
pPlinkHet.lang         = params.Rscript.value

pPlinkHWE = proc_factory(
	desc = "Hardy-Weinberg Equilibrium report and filtering.",
	config = Diot(annotate = """
	@name:
		pPlinkHWE
	@description:
		Hardy-Weinberg Equilibrium report and filtering.
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outdir:dir`: The output directory. Default: `{{i.indir | fn}}.hwe`
			- `.hwe`: The HWE report by `plink`
			- `.hardy.fail`: The SNPs fail HWE test
	@args:
		`plink`: The path to plink.
		`cutoff`: The HWE p-value cutoff. Default: `1e-5`
		`plot`: Whether plot the distribution of the HWE p-values? Default: `True`
		`devpars`: The device parameters for the plot. Default: `Diot(res=300, width=2000, height=2000)`
	"""))
pPlinkHWE.input        = 'indir:dir'
pPlinkHWE.output       = 'outdir:dir:{{i.indir | fn}}.hwe'
pPlinkHWE.args.plink   = params.plink.value
pPlinkHWE.args.cutoff  = 1e-5
pPlinkHWE.args.mingt   = .05
pPlinkHWE.args.plot    = True
pPlinkHWE.args.devpars = Diot(res=300, width=2000, height=2000)
pPlinkHWE.lang         = params.Rscript.value

pPlinkIBD = proc_factory(
	desc = "Estimate the identity by descent (IBD)",
	config = Diot(annotate = """
	@name:
		pPlinkIBD
	@description:
		Estimate the degree of recent shared ancestry individual pairs,
		the identity by descent (IBD)
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outdir:dir`: The output directory. Default: `{{i.indir | fn}}.ibd`
			- `.genome`: The original genome report from `plink`
			- `.ibd.png`: The heatmap of PI_HAT
	@args:
		`plink`: The path to plink.
		`indep`: To give a concrete example: the command above that specifies 50 5 0.2 would a) consider a window of 50 SNPs, b) calculate LD between each pair of SNPs in the window, b) remove one of a pair of SNPs if the LD is greater than 0.5, c) shift the window 5 SNPs forward and repeat the procedure.
		`pihat`: The PI_HAT cutoff. Default: 0.1875 (see: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5007749/)
		`plot` : Whether plot the PI_HAT heatmap? Default: `True`
		`devpars`: The device parameters for the plot. Default: `Diot(res=300, width=2200, height=1600)`
		`samid`: Sample ids on the heatmap. Default: `iid`
			- Could also be `fid` or `fid<sep>iid`, or an R function: `function(fid, iid)`
		`anno` : The annotation file for the samples. Names must match the ones that are transformed by `args.samid`. Default: `''`
	"""))
pPlinkIBD.input        = 'indir:dir'
pPlinkIBD.output       = 'outdir:dir:{{i.indir | fn}}.ibd'
pPlinkIBD.args.plink   = params.plink.value
pPlinkIBD.args.highld  = params.highld.value
pPlinkIBD.args.samid   = 'iid' # fid or a function (fid, iid)
pPlinkIBD.args.indep   = [50, 5, .2]
pPlinkIBD.args.pihat   = 0.1875 # ref: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5007749/
pPlinkIBD.args.plot    = True
pPlinkIBD.args.anno    = ''
pPlinkIBD.args.seed    = None
pPlinkIBD.args.devpars = Diot(res=300, width=2000, height=2000)
pPlinkIBD.lang         = params.Rscript.value

pPlinkRemove = proc_factory(
	desc = "Remove failed samples and SNPs",
	config = Diot(annotate = """
	@name:
		pPlinkRemove
	@description:
		Remove failed samples and/or SNPs
		The samples/SNPs to be removed should be generated by one of:
		`pPlinkHet`, `pPlinkHWE`, `pPlinkIBD` or `pPlinkMiss`
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
		`pdir:dir` : The output directory from one of the processes listed in description
			- It could also be the `.fail` file generated by those processes
	@output:
		`outdir:dir`: The output directory containing the `.bed/.bim/.fam` after filtering.
	@args:
		`plink`: The path to plink.
	"""))
pPlinkRemove.input        = 'indir:dir, pdir:dir'
pPlinkRemove.output       = 'outdir:dir:{{i.indir | fn}}'
pPlinkRemove.args.plink   = params.plink.value
pPlinkRemove.envs.rimport = rimport
pPlinkRemove.lang         = params.Rscript.value

pPlink2Vcf = proc_factory(
	desc = "Convert plink binary files to VCF file.",
	config = Diot(annotate = """
	@name:
		pPlink2Vcf
	@description:
		Convert plink binaries into VCF file.
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outfile:file`: The output vcf file.
	@args:
		`plink`: The path to plink.
		`gz`   : Whether bgzip the output vcf file. Default: `False`
		`samid`: What to use as sample ID. Default: `both`
			- `both`: use `<FID>_<IID>` as sample id
			- `fid` : use `<FID>` as sample id
			- `iid` : use `<IID>` as sample id
	"""))
pPlink2Vcf.input       = 'indir:dir'
pPlink2Vcf.output      = 'outfile:file:{{i.indir | bn}}.vcf{% if args.gz %}.gz{% endif %}'
pPlink2Vcf.args.plink  = params.plink.value
pPlink2Vcf.args.gz     = False
pPlink2Vcf.args.samid  = 'both' # fid, iid
pPlink2Vcf.args.chroms = {"23": "X", "24": "Y", "25": "XY", "26": "M"}
pPlink2Vcf.lang        = params.python.value

pPlink2GTMat = proc_factory(
	desc = "Convert plink binary files to genotype matrix",
	config = Diot(annotate = """
	@name:
		pPlink2GTMat
	@description:
		Convert plink binaries into genotype matrix.
	@input:
		`indir:dir`: The input directory containing .bed/.bim/.fam files
	@output:
		`outfile:file`: The output genotype matrix file.
	@args:
		`plink`: The path to plink.
		`samid`: What to use as sample ID. Default: `both`
			- `both`: use `<FID>_<IID>` as sample id
			- `fid` : use `<FID>` as sample id
			- `iid` : use `<IID>` as sample id
	"""))
pPlink2GTMat.input       = 'indir:dir'
pPlink2GTMat.output      = 'outfile:file:{{i.indir | bn}}.gtmat.txt'
pPlink2GTMat.args.plink  = params.plink.value
pPlink2GTMat.args.samid  = 'both' # fid, iid
pPlink2GTMat.args.addchr = True
pPlink2GTMat.args.snpid  = '{chr}_{pos}_{rs}_{ref}_{alt}' # or raw
pPlink2GTMat.args.chroms = {"23": "X", "24": "Y", "25": "XY", "26": "M"}
pPlink2GTMat.args.nors   = "NOVEL"
pPlink2GTMat.lang        = params.python.value

pPlinkPCA = proc_factory(
	desc = "Perform PCA on genotype data and covariates.",
	config = Diot(annotate = """
	@name:
		pPlinkPCA
	@description:
		Do PCA on genotype data with PLINK
	@input:
		`indir`: The input directory with .bed/.bim/.fam files
	@output:
		`outfile:file`: The output file of selected PCs, Default: `{{i.indir | fn}}.plinkPCA/{{i.indir | fn}}.pcs.txt`
		`outdir:dir`: The output directory with output file and plots. Default: `{{i.indir | fn}}.plinkPCA`
	@args:
		`plink`: The path to `plink`, Default: `<params.plink>`
		`samid`: Which IDs to report in results, Default: `both`
			- `both`: Both family ID and individual ID connected with `_`
			- `iid`:  Individual ID
			- `fid`:  Family ID
		`nthread`: # threads to use, Default: `False`
			- `False`: Don't put `--threads` in plink command
		`indep`: `indep` used to prune LD SNPs. Default: `[50, 5, .2]`
		`highld`: High LD regions. Default: `<params.highld>`
		`params`: Other parameters for `plink --pca`. Default: `Diot(mind = .95)`
		`select`: Select first PCs in the output file. Default: `0.2`
			- `select < 1`: select PCs with contribution greater than `select`
			- `select >=1`: select first `select` PCs
		`plots` : Output plots. Default: `Diot(scree = Diot(ncp = 20))`
		`devpars`: The parameters for ploting device. Default: `Diot(height = 2000, width = 2000, res = 300)`
	"""))
pPlinkPCA.input    = 'indir:dir'
pPlinkPCA.output   = [
	'outfile:file:{{i.indir | fn}}.plinkPCA/{{i.indir | fn}}.pcs.txt',
	'outdir:dir:{{i.indir | fn}}.plinkPCA'
]
pPlinkPCA.args.plink   = params.plink.value
pPlinkPCA.args.samid   = 'both' # fid, iid
pPlinkPCA.args.nthread = False
pPlinkPCA.args.indep   = [50, 5, .2] # used to prune LD SNPs
pPlinkPCA.args.highld  = params.highld.value
pPlinkPCA.args.params  = Diot(mind = .95)
pPlinkPCA.args.select  = .2
pPlinkPCA.args.plots   = Diot(
	scree = Diot(ncp = 20),
	# rownames of anno should be consistent with `args.samid`
	pairs = Diot(anno = '', ncp = 4, params = Diot(upper = Diot(continuous = 'density')), ggs = Diot(theme = {"axis.text.x": "r:ggplot2::element_text(angle = 60, hjust = 1)" })),
	# more to add
)
pPlinkPCA.args.devpars = Diot(height = 2000, width = 2000, res = 300)
pPlinkPCA.envs.rimport = rimport
pPlinkPCA.lang         = params.Rscript.value

pPlinkSimulate = proc_factory(
	desc = "Simulate a set of SNPs",
	config = Diot(annotate = """
	@name:
		pPlinkSimulate
	"""))
pPlinkSimulate.input        = 'seed'
pPlinkSimulate.output       = 'outdir:dir:simsnps.{{i.seed if isinstance(i.seed, int) else "noseed"}}.plink'
pPlinkSimulate.args.plink   = params.plink.value
pPlinkSimulate.args.ncases  = 1000
pPlinkSimulate.args.nctrls  = 1000
pPlinkSimulate.args.nsnps   = 100
pPlinkSimulate.args.label   = 'SNP'
pPlinkSimulate.args.dprev   = .01
pPlinkSimulate.args.minfreq = 0
pPlinkSimulate.args.maxfreq = 1
pPlinkSimulate.args.hetodds = 1
pPlinkSimulate.args.homodds = 1
pPlinkSimulate.args.params  = Diot()
pPlinkSimulate.lang         = params.python.value
