from autofit import conf
from autofit.tools import path_util
from autocti.data import util
from autocti.charge_injection import ci_frame
from autocti.charge_injection import ci_data
from autocti.charge_injection import ci_pattern

from autocti.model import arctic_settings

from multiprocessing import Pool

import os
import sys

### NOTE - if you have not already, complete the setup in 'workspace/runners/cosma/setup' before continuing with this
### cosma pipeline script.

# Welcome to the Cosma pipeline runner. Hopefully, you're familiar with runners at this point, and have been using them
# with PyAutoCTI to model CTI on your laptop. If not, I'd recommend you get used to doing that, before trying to
# run PyAutoCTI on a super-computer. You need some familiarity with the software and before trying to model a large
# amount of charge injection imaging on a supercomputer!

# If you are ready, then let me take you through the Cosma runner. It is remarkably similar to the ordinary pipeline
# runners you're used to, however it makes a few changes for running jobs on cosma:

# 1) The data path is over-written to the path '/cosma5/data/durham/cosma_username/autocti/data' as opposed to the
#    workspace. As we saw in the setup, on cosma we don't store our data in our workspace.

# 2) The output path is over-written to the path '/cosma5/data/durham/cosma_username/autocti/output' as opposed to
#    the workspace. This is for the same reason as the data.

# Given your username is where your data is stored, you'll need to put your cosma username here.
cosma_username = 'pdtw24'

# The cosma path where the data and output are stored.
cosma_path = '/cosma5/data/autocti/'

# Get the relative path to the config files and output folder in our workspace.
workspace_path = '{}/../'.format(os.path.dirname(os.path.realpath(__file__)))

# Use this path to explicitly set the config path, and override the output path with the Cosma path.
conf.instance = conf.Config(config_path=workspace_path + 'config', output_path=cosma_path + 'output')

# Lets take a look at a Cosma batch script, which can be found at 'workspace/runners/cosma/batch/pipeline_runner_cosma'.
# When we submit a PyAutoLens job to Cosma, we submit a 'batch' of jobs, whereby each job will run on one CPU of Cosma.
# Thus, if our lens sample contains, lets say, 4 lenses, we'd submit 4 jobs at the same time where each job applies
# our pipeline to each image.

# The fifth line of this batch script - '#SBATCH --array=1-4' is what species this. Its telling Cosma we're going to
# run 4 jobs, and the id's of those jobs will be numbered from 1 to 4. Infact, these ids are passed to this runner,
# and we'll use them to ensure that each jobs loads a different image. Lets get the cosma array id for our job.
batch_id = int(sys.argv[1])
ci_data_resolution = str(sys.argv[2]).strip() # The resolution of the image.

# For a given COSMA run, we will assume all of the charge injection data-sets have identical properties. That is, they
# all have the same dimensions (shape), charge injection regions (ci_regions), charge injection normalizations
# (normalizations), charge injection patterns (patterns) and frame geometries defining the direction clocking and CTI
# (frame_geometry).

# If you need to change these value for a COSMA run, I recommend making a different runner script.

ci_data_type = 'ci_images_non_uniform' # Charge injection data consisting of 2 images with uniform injections.

ci_data_model = []
ci_data_model.append('') # Task number beings at 1, so keep index 0 blank
ci_data_model.append('serial_x3_species_model_7') # Index 6

# Create the path where the data will be loaded from, which in this case is
# '/workspace/data/ci_images_uniform/parallel_x2_species/high_res/'
ci_data_path = path_util.make_and_return_path_from_path_and_folder_names(
    path=cosma_path, folder_names=['data', ci_data_type, ci_data_model[batch_id], ci_data_resolution])

# The image resolution determines the image shape and frame geometry.
if ci_data_resolution in 'high_res':
    shape = (2316, 2119)
    frame_geometry = ci_frame.FrameGeometry(corner=(0.0, 0.0),
                                            parallel_overscan=ci_frame.Region((2296, 2316, 51, 2099)),
                                            serial_prescan=ci_frame.Region((0, 2316, 0, 51)),
                                            serial_overscan=ci_frame.Region((0, 2296, 2099, 2119)))
elif ci_data_resolution in 'mid_res':
    shape = (1158, 2119)
    frame_geometry = ci_frame.FrameGeometry(corner=(0.0, 0.0),
                                            parallel_overscan=ci_frame.Region((1138, 1158, 51, 2099)),
                                            serial_prescan=ci_frame.Region((0, 1158, 0, 51)),
                                            serial_overscan=ci_frame.Region((0, 1138, 2099, 2119)))
elif ci_data_resolution in 'low_res':
    shape = (579, 2119)
    frame_geometry = ci_frame.FrameGeometry(corner=(0.0, 0.0),
                                            parallel_overscan=ci_frame.Region((559, 579, 51, 2099)),
                                            serial_prescan=ci_frame.Region((0, 579, 0, 51)),
                                            serial_overscan=ci_frame.Region((0, 559, 2099, 2119)))

# The charge injection regions on the CCD, which in this case is 7 equally spaced rectangular blocks.
ci_regions = [(0, shape[0], 51, shape[1]-20)]

# The normalization of the charge injection for each image.
normalizations=[100.0, 500.0, 1000.0, 5000.0, 10000.0, 25000.0, 50000.0, 84700.0]

# Create the charge injection pattern objects used for this pipeline.
patterns = ci_pattern.uniform_from_lists(normalizations=normalizations, regions=ci_regions)

ci_datas = []

for index, normalization in enumerate(normalizations):

    ci_datas.append(ci_data.ci_data_from_fits(
                 frame_geometry=frame_geometry, ci_pattern=patterns[index],
                 image_path=ci_data_path + 'image_' + str(int(normalization)) + '.fits',
                 ci_pre_cti_path=ci_data_path +'ci_pre_cti_' + str(int(normalization)) + '.fits',
                 noise_map_from_single_value=4.0))

serial_cti_settings = arctic_settings.Settings(well_depth=84700, niter=1, express=2, n_levels=2000,
                                               charge_injection_mode=False, readout_offset=0)
cti_settings = arctic_settings.ArcticSettings(serial=serial_cti_settings)

# Running a pipeline is exactly the same as we're used to. We import it, make it, and run it, noting that we can
# use the lens_name's to ensure each job outputs its results to a different directory.

from workspace_jam.pipelines import serial_x3_species_fix

pipeline = serial_x3_species_fix.make_pipeline(phase_folders=[ci_data_type, ci_data_model[batch_id], ci_data_resolution])

pipeline.run(ci_datas=ci_datas, cti_settings=cti_settings, pool=Pool(processes=8))

# Finally, its worth us going through a batch script in detail, line by line, as you may we need to change different
# parts of this script to use different runners. Therefore, checkout the 'doc' file in the batch folder.