<commit_before>""" All about Artificial star tests """
from __future__ import print_function
import os
import numpy as np

import logging
logger = logging.getLogger()

from scipy.interpolate import interp1d
from subprocess import PIPE, Popen

from .. import trilegal
from .. import io
from .. import astronomy_utils
__all__ = ['ast_correct_trilegal_sim', 'ASTs']


def ast_correct_trilegal_sim(sgal, fake_file=None, outfile=None,
                             overwrite=False, spread_outfile=None,
                             leo_method=False, spread_outfile2=None,
                             asts_obj=None):
    '''
    correct trilegal simulation with artificial star tests.
    options to write the a copy of the trilegal simulation with the corrected
    mag columns, and also just the color, mag into a file (spread_too).

    ARGS:
    sgal: simgalaxy instance with convert_mag method already called.
    fake_file: string - matchfake file(s), or an artifical_star_tests instance(s).
    outfile: string - ast_file to write
    overwrite: overwrite output files, is also passed to write_spread
    spread_too: call write_spread flag
    spread_outfile: outfile for write_spread

    RETURNS:
    adds corrected mags to sgal.data.data_array and updates sgal.data.key_dict

    if fake_file is an a list of opt and ir, will do the corrections twice.
    fake_file as artificial_star_instance won't work if leo_method=True.
    '''
    ir_too = False

    def leo_ast_correction(leo_code, fake_file, sim_file, spread_outfile):
        # <hit enter>
        EOF = os.path.join(os.environ['PYTHONPATH'], 'EOF')

        cmd = '%s << %s\n' % (leo_code, EOF)
        cmd += '\n'.join((fake_file, sim_file, spread_outfile, EOF))

        logger.debug("%s + %s -> %s" % (sim_file, fake_file, spread_outfile))
        logger.debug('Running %s...' % os.path.split(leo_code)[1])
        p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE,
                  close_fds=True)

        stdout, stderr = (p.stdout, p.stderr)
        p.wait()
        return

    def add_cols_leo_method(sgal, spread_outfile):
        spout = io.readfile(spread_outfile)
        ast_filts = [s for s in spout.dtype.names if '_cor' in s]
        new_cols = {}
        for ast_filt in ast_filts:
            new_cols[ast_filt] = spout[ast_filt]
            print(len(spout[ast_filt]))
        print(len(sgal.mag2))
        print(sgal.add_data(**new_cols))

    #assert fake_file is not None and asts_obj is not None, \
    #    'ast_correct_trilegal_sim: fake_file now needs to be passed'

    if leo_method is True:
        print('ARE YOU SURE YOU WANT TO USE THIS METHOD!?')
        # this method tosses model stars where there are no ast corrections
        # which is fine for completeness < .50 but not for completeness > .50!
        # so don't use it!! It's also super fucking slow.
        if spread_outfile is None:
            raise ValueError('need spread_outfile set for Leo AST method')

        print("completeness using Leo's method")

        sim_file = os.path.join(sgal.base, sgal.name)

        # special version of spread_angst made for wfc3snap photsys!
        # THIS IS BAD CODING HARD LINK?!
        leo_code = '/home/rosenfield/research/TP-AGBcalib/SNAP/models/spread_angst'

        if type(fake_file) == list:
            assert spread_outfile2 is not None, \
                'need another string for spread file name with four filters'
            fake_file_ir = fake_file[1]
            fake_file = fake_file[0]
            ir_too = True

        leo_ast_correction(leo_code, fake_file, sim_file, spread_outfile)
        add_cols_leo_method(sgal, spread_outfile)
        if ir_too is True:
            # run the out ast corrections through a version of spread_angst to
            # correct the ir.
            leo_code += '_ir'
            leo_ast_correction(leo_code, fake_file_ir, spread_outfile, spread_outfile2)
            add_cols_leo_method(sgal, spread_outfile2)
            # no need to keep the spread_outfile now that we have a new one.
            os.remove(spread_outfile)
    else:
        if type(fake_file) is str:
            fake_files = [fake_file]
        else:
            fake_files = fake_file

        sgal.fake_files = fake_files

        if asts_obj is None:
            asts_obj = [ASTs(fake_file) for fake_file in fake_files]
        elif type(asts_obj) is str:
            asts_obj = [asts_obj]

        for asts in asts_obj:
            # sgal.filter1 or sgal.mag1 doesn't have to match the asts.filter1
            # in fact, it won't if this is done on both opt and nir data.
            mag1 = sgal.data.get_col(asts.filter1)
            mag2 = sgal.data.get_col(asts.filter2)

            cor_mag1, cor_mag2 = asts.ast_correction(mag1, mag2,
                                                     **{'binsize': 0.2})
            new_cols = {'%s_cor' % asts.filter1: cor_mag1,
                        '%s_cor' % asts.filter2: cor_mag2}
            sgal.add_data(**new_cols)

        if outfile is not None:
            if overwrite is True or not os.path.isfile(outfile):
                trilegal.write_trilegal_sim(sgal, outfile)
            else:
                logger.warning('%s exists, not overwriting' % outfile)

        if spread_outfile is not None:
            trilegal.write_spread(sgal, outfile=spread_outfile, overwrite=overwrite)


class ASTs(object):
    '''
    class for reading and using artificial stars.
    If *filename* is not a matchfake file, MUST create a new method to read in
    artificial stars.
    for load_fake:
    mag1 is assumed to be mag1in
    mag2 is assumed to be mag2in
    mag1diff is assumed to be mag1in-mag1out
    mag2diff is assumed to be mag2in-mag2out
    filename is assumed as matchfake:
    PID_TARGET_FILTER1_FILTER2_... or TARGET_FILTER1_FILTER2_
    this is how attributes target, filter1, and filter2 are assigned.

    '''
    def __init__(self, filename):
        self.base, self.name = os.path.split(filename)
        try:
            __, self.target, self.filter1, self.filter2, _ = \
                self.name.split('_')
        except:
            try:
                self.target, self.filter1, filter2 = self.name.split('_')
                self.filter2 = filter2.split('.')[0]
            except:
                try:
                    __, self.target, __, self.filter1, self.filter2, _ = \
                        self.name.split('_')
                except:
                    pass
        self.load_fake(filename)

    def recovered(self, threshold=9.99):
        '''
        find indicies of stars with magdiff < threshold

        ARGS:
        threshold: [9.99] magin - magout threshold for recovery

        RETURNS:
        self.rec: recovered stars in both filters
        rec1, rec2: recovered stars in filter1, filter2
        '''
        rec1, = np.nonzero(self.mag1diff < threshold)
        rec2, = np.nonzero(self.mag2diff < threshold)
        self.rec = list(set(rec1) & set(rec2))
        return rec1, rec2

    def ast_hess(self, binsize=0.1, yattr='mag2diff', hess_kw={}):
        '''
        make hess grid
        '''
        self.colordiff = self.mag1diff - self.mag2diff
        mag = self.__getattribute__(yattr)
        self.hess = astronomy_utils.hess(self.colordiff, mag, binsize,
                                         **hess_kw)

    def load_fake(self, filename):
        '''read MATCH fake file into attributes'''
        names = ['mag1', 'mag2', 'mag1diff', 'mag2diff']
        self.data = np.genfromtxt(filename, names=names)
        # unpack into attribues
        for name in names:
            self.__setattr__(name, self.data[name])

    def bin_asts(self, binsize=0.2, bins=None):
        '''
        bin the artificial star tests

        ARGS:
        bins: bins for the asts
        binsize: width of bins for the asts

        RETURNS:<commit_msg>ast_hess ... bad function name<commit_after>""" All about Artificial star tests """
from __future__ import print_function
import os
import numpy as np

import logging
logger = logging.getLogger()

from scipy.interpolate import interp1d
from subprocess import PIPE, Popen

from .. import trilegal
from .. import io
from .. import astronomy_utils
__all__ = ['ast_correct_trilegal_sim', 'ASTs']


def ast_correct_trilegal_sim(sgal, fake_file=None, outfile=None,
                             overwrite=False, spread_outfile=None,
                             leo_method=False, spread_outfile2=None,
                             asts_obj=None):
    '''
    correct trilegal simulation with artificial star tests.
    options to write the a copy of the trilegal simulation with the corrected
    mag columns, and also just the color, mag into a file (spread_too).

    ARGS:
    sgal: simgalaxy instance with convert_mag method already called.
    fake_file: string - matchfake file(s), or an artifical_star_tests instance(s).
    outfile: string - ast_file to write
    overwrite: overwrite output files, is also passed to write_spread
    spread_too: call write_spread flag
    spread_outfile: outfile for write_spread

    RETURNS:
    adds corrected mags to sgal.data.data_array and updates sgal.data.key_dict

    if fake_file is an a list of opt and ir, will do the corrections twice.
    fake_file as artificial_star_instance won't work if leo_method=True.
    '''
    ir_too = False

    def leo_ast_correction(leo_code, fake_file, sim_file, spread_outfile):
        # <hit enter>
        EOF = os.path.join(os.environ['PYTHONPATH'], 'EOF')

        cmd = '%s << %s\n' % (leo_code, EOF)
        cmd += '\n'.join((fake_file, sim_file, spread_outfile, EOF))

        logger.debug("%s + %s -> %s" % (sim_file, fake_file, spread_outfile))
        logger.debug('Running %s...' % os.path.split(leo_code)[1])
        p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE,
                  close_fds=True)

        stdout, stderr = (p.stdout, p.stderr)
        p.wait()
        return

    def add_cols_leo_method(sgal, spread_outfile):
        spout = io.readfile(spread_outfile)
        ast_filts = [s for s in spout.dtype.names if '_cor' in s]
        new_cols = {}
        for ast_filt in ast_filts:
            new_cols[ast_filt] = spout[ast_filt]
            print(len(spout[ast_filt]))
        print(len(sgal.mag2))
        print(sgal.add_data(**new_cols))

    #assert fake_file is not None and asts_obj is not None, \
    #    'ast_correct_trilegal_sim: fake_file now needs to be passed'

    if leo_method is True:
        print('ARE YOU SURE YOU WANT TO USE THIS METHOD!?')
        # this method tosses model stars where there are no ast corrections
        # which is fine for completeness < .50 but not for completeness > .50!
        # so don't use it!! It's also super fucking slow.
        if spread_outfile is None:
            raise ValueError('need spread_outfile set for Leo AST method')

        print("completeness using Leo's method")

        sim_file = os.path.join(sgal.base, sgal.name)

        # special version of spread_angst made for wfc3snap photsys!
        # THIS IS BAD CODING HARD LINK?!
        leo_code = '/home/rosenfield/research/TP-AGBcalib/SNAP/models/spread_angst'

        if type(fake_file) == list:
            assert spread_outfile2 is not None, \
                'need another string for spread file name with four filters'
            fake_file_ir = fake_file[1]
            fake_file = fake_file[0]
            ir_too = True

        leo_ast_correction(leo_code, fake_file, sim_file, spread_outfile)
        add_cols_leo_method(sgal, spread_outfile)
        if ir_too is True:
            # run the out ast corrections through a version of spread_angst to
            # correct the ir.
            leo_code += '_ir'
            leo_ast_correction(leo_code, fake_file_ir, spread_outfile, spread_outfile2)
            add_cols_leo_method(sgal, spread_outfile2)
            # no need to keep the spread_outfile now that we have a new one.
            os.remove(spread_outfile)
    else:
        if type(fake_file) is str:
            fake_files = [fake_file]
        else:
            fake_files = fake_file

        sgal.fake_files = fake_files

        if asts_obj is None:
            asts_obj = [ASTs(fake_file) for fake_file in fake_files]
        elif type(asts_obj) is str:
            asts_obj = [asts_obj]

        for asts in asts_obj:
            # sgal.filter1 or sgal.mag1 doesn't have to match the asts.filter1
            # in fact, it won't if this is done on both opt and nir data.
            mag1 = sgal.data.get_col(asts.filter1)
            mag2 = sgal.data.get_col(asts.filter2)

            cor_mag1, cor_mag2 = asts.ast_correction(mag1, mag2,
                                                     **{'binsize': 0.2})
            new_cols = {'%s_cor' % asts.filter1: cor_mag1,
                        '%s_cor' % asts.filter2: cor_mag2}
            sgal.add_data(**new_cols)

        if outfile is not None:
            if overwrite is True or not os.path.isfile(outfile):
                trilegal.write_trilegal_sim(sgal, outfile)
            else:
                logger.warning('%s exists, not overwriting' % outfile)

        if spread_outfile is not None:
            trilegal.write_spread(sgal, outfile=spread_outfile, overwrite=overwrite)


class ASTs(object):
    '''
    class for reading and using artificial stars.
    If *filename* is not a matchfake file, MUST create a new method to read in
    artificial stars.
    for load_fake:
    mag1 is assumed to be mag1in
    mag2 is assumed to be mag2in
    mag1diff is assumed to be mag1in-mag1out
    mag2diff is assumed to be mag2in-mag2out
    filename is assumed as matchfake:
    PID_TARGET_FILTER1_FILTER2_... or TARGET_FILTER1_FILTER2_
    this is how attributes target, filter1, and filter2 are assigned.

    '''
    def __init__(self, filename):
        self.base, self.name = os.path.split(filename)
        try:
            __, self.target, self.filter1, self.filter2, _ = \
                self.name.split('_')
        except:
            try:
                self.target, self.filter1, filter2 = self.name.split('_')
                self.filter2 = filter2.split('.')[0]
            except:
                try:
                    __, self.target, __, self.filter1, self.filter2, _ = \
                        self.name.split('_')
                except:
                    pass
        self.load_fake(filename)

    def recovered(self, threshold=9.99):
        '''
        find indicies of stars with magdiff < threshold

        ARGS:
        threshold: [9.99] magin - magout threshold for recovery

        RETURNS:
        self.rec: recovered stars in both filters
        rec1, rec2: recovered stars in filter1, filter2
        '''
        rec1, = np.nonzero(self.mag1diff < threshold)
        rec2, = np.nonzero(self.mag2diff < threshold)
        self.rec = list(set(rec1) & set(rec2))
        return rec1, rec2

    def make_hess(self, binsize=0.1, yattr='mag2diff', hess_kw={}):
        '''
        make hess grid
        '''
        self.colordiff = self.mag1diff - self.mag2diff
        mag = self.__getattribute__(yattr)
        self.hess = astronomy_utils.hess(self.colordiff, mag, binsize,
                                         **hess_kw)

    def load_fake(self, filename):
        '''read MATCH fake file into attributes'''
        names = ['mag1', 'mag2', 'mag1diff', 'mag2diff']
        self.data = np.genfromtxt(filename, names=names)
        # unpack into attribues
        for name in names:
            self.__setattr__(name, self.data[name])

    def bin_asts(self, binsize=0.2, bins=None):
        '''
        bin the artificial star tests

        ARGS:
        bins: bins for the asts
        binsize: width of bins for the asts

        RETURNS:
{'fraction_non_alphanumeric': 0.08664556962025316, 'fraction_numerical': 0.010759493670886076, 'mean_word_length': 2.1551517571884986, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4776340', 'n_tokens_mistral': 5000, 'n_tokens_neox': 4794, 'n_words': 1394}