#!/bin/python

configfile: "config/config.json"

import numpy as np 
import numpy as np 
import pandas as pd
import sys
sys.path.append('src/')
shell.prefix("set -o pipefail; ")

# Define the data directory
ANNO_DIR = config['ANNO_DIR']
DATA_DIR = config['DATA_DIR']

# ------------ 0. Preliminaries for plotting ----------------- # 
rule plotting_prelims:
  input:
    rules.count_geodist_categories_subset_all.input,
    rules.count_geodist_all.input


# ------------- 1. Full Population Plot ---------------------# 
rule plot_full_pops:
  input:
    'data/freq/new_1kg_nyc_hg38_filt_chr22.biallelic_snps.pops.freq.txt.gz'
  output:
    fig1='plots/figure1/fig1.pdf'
  shell:
    """
    cd doc/
    python3 -W ignore plot_full_pop.py
    cd ..
    """

# --------------- 2. Overall Figures ----------------------- # 
rule plot_overall_figures:
  input:
    rules.plotting_prelims.input
  output:
    fig3B = 'plots/figure3/fig3B.pdf',
    fig3C = 'plots/figure3/fig3C.pdf',
    figS1A = 'plots/figureS1/figS1A.pdf',
    figS1B = 'plots/figureS1/figS1B.pdf',
    fig7 = 'plots/figure7/fig7.pdf'
  shell:
    """
      cd doc/
      python3 -W ignore plot_overall.py
      cd ..
    """


# --------------- 3. Genotyping Array Figures ---------------- #
rule plot_genotyping_array_figures:
  input:
    rules.plotting_prelims.input
  output:
    array_gallery = 'plots/figure6/fig6.pdf'
  shell:
    """
      cd doc/
      python3 -W ignore plot_arrays.py
      cd ..
    """

# --------------- 4. SGDP Figures Generation --------------- # 
rule plot_sgdp_figures:
  input:
    rules.plotting_prelims.input
  output:
    sgdp_gallery = 'plots/figure5/fig5.pdf'
  shell:
    """
      cd doc/
      python3 -W ignore plot_sgdp.py
      cd ..
    """
    
#--------------- 5. Theory Figures Generation ------------- #
rule plot_theory_figures:
  output:
    theory_geodist_gallery = 'plots/figure4/fig4A.pdf',
    theory_envelopes = 'plots/figure4/fig4BC.pdf'
  shell:
    """
      cd doc/
      python3 -W ignore plot_jsfs_theory.py
      python3 -W ignore heuristic_figure.py
      cd ..
    """
    
rule gen_all_plots:
  """
    Meta-Rule to generate all of the main figures! 
  """
  input:
    rules.plot_full_pops.output,
    rules.plot_overall_figures.output,
    rules.plot_genotyping_array_figures.output,
    rules.plot_sgdp_figures.output,
    rules.plot_theory_figures.output