#!/usr/bin/env python
"""
Implements LDpred, an approximate Gibbs sampler that calculate posterior means of effects, conditional on LD information.
The method requires the user to have generated a coordinated dataset using coord_genotypes.py

Usage:
ldpred --coord=COORD_DATA_FILE  --ld_radius=LD_RADIUS   --local_ld_file_prefix=LD_FILE_NAME  --PS=FRACTIONS_CAUSAL
                          --N=SAMPLE_SIZE  --out=OUTPUT_FILE_PREFIX  [ --num_iter=NUM_ITER  --H2=HERTIABILITY  --gm_ld_radius=GEN_MAP_RADIUS]

 - COORD_DATA_FILE: The HDF4 file obtained by running the coord_genotypes.py

 - LD_RADIUS: An integer number which denotes the number of SNPs on each side of the focal SNP for which LD should be adjusted.
              A value corresponding M/3000, where M is the number of SNPs in the genome is recommended.

 - LD_FILE_NAME: A path and filename prefix for the LD file.  If it doesn't exist, it will be generated.  This can take up to several hours,
                 depending on LD radius number of SNPs, etc.  If it does exits, that file will be used.

 - FRACTIONS_CAUSAL: A list of comma separated (without space) values between 1 and 0, excluding 0.  1 corresponds to the infinitesimal model and will yield results
                     similar to LDpred-inf.  Default is --PS=1,0.3,0.1,0.03,0.01,0.003,0.001,0.0003,0.0001

 - N: This is the sample size which LDpred assumes was used to calculate the GWAS summary statistics.

 - OUTPUT_FILE_PREFIX:  The prefix of output file.

 - NUM_ITER (optional): The number of iterations used by the Gibbs sampler. The default is 60, and burn-in is fixed to 5.

 - HERTIABILITY (optional): The heritability assumed by LDpred.  By default it estimates the heritability from the GWAS summary statistics.

 - GEN_MAP_RADIUS (optional):  If this option is set, then a genetic map will be used to calculate LD-radius.  A value around 1 is arguably reasonable.

 2015 (c) Bjarni J Vilhjalmsson: bjarni.vilhjalmsson@gmail.com

 """

import cPickle
import getopt
import gzip
import os
import sys
import time

from scipy import stats
import h5py

import LDpred_inf
import itertools as it
import ld
import scipy as sp

from coord_genotypes import chromosomes_list
import util

__version__ = '0.9.1'

def parse_parameters():
    """
    Parse the parameters into a dict, etc.
    """

    long_options_list = ['coord=', 'ld_radius=', 'local_ld_file_prefix=', 'PS=', 'out=', 'N=',
                         'num_iter=', 'H2=', 'gm_ld_radius=', 'h', 'help']

    p_dict = {'coord':None, 'ld_radius':None, 'local_ld_file_prefix':None, 'PS':[1, 0.3, 0.1, 0.03, 0.01, 0.003, 0.001], 'out':None,
              'N':None, 'num_iter': 60, 'H2':None, 'gm':None, 'gm_ld_radius':None}

    if len(sys.argv) > 1:
        try:
            opts, args = getopt.getopt(sys.argv[1:], "h", long_options_list)

        except:
            print "Some problems with parameters.  Please read the usage documentation carefully."
            print "Use the -h option for usage information."
            sys.exit(2)

        for opt, arg in opts:
            if opt == "-h" or opt == "--h" or opt == '--help':
                print __doc__
                sys.exit(0)
            elif opt == "--coord": p_dict['coord'] = arg
            elif opt == "--ld_radius": p_dict['ld_radius'] = int(arg)
            elif opt == "--local_ld_file_prefix": p_dict['local_ld_file_prefix'] = arg
            elif opt == "--PS": p_dict['PS'] = map(float, arg.split(','))
            elif opt == "--out": p_dict['out'] = arg
            elif opt == "--N": p_dict['N'] = int(arg)
            elif opt == "--num_iter": p_dict['num_iter'] = int(arg)
            elif opt == "--H2": p_dict['H2'] = float(arg)
            elif opt == "--gm_ld_radius": p_dict['gm_ld_radius'] = float(arg)
            else:
                print "Unkown option:", opt
                print "Use -h option for usage information."
                sys.exit(2)
    else:
        print __doc__
        sys.exit(0)
    return p_dict

def ldpred_genomewide(data_file=None, ld_radius=None, ld_dict=None, out_file_prefix=None, ps=None,
               n=None, h2=None, num_iter=None, verbose=False, zero_jump_prob=0.05, burn_in=5):
    """
    Calculate LDpred for a genome
    """

    df = h5py.File(data_file, 'r')
    has_phenotypes = False
    if 'y' in df.keys():
        'Validation phenotypes found.'
        y = df['y'][...]  # Phenotype
        num_individs = len(y)
        risk_scores_pval_derived = sp.zeros(num_individs)
        has_phenotypes = True

    ld_scores_dict = ld_dict['ld_scores_dict']
    chrom_ld_dict = ld_dict['chrom_ld_dict']
    chrom_ref_ld_mats = ld_dict['chrom_ref_ld_mats']

    print 'Applying LDpred with LD radius: %d' % ld_radius
    results_dict = {}
    num_snps = 0
    sum_beta2s = 0
    cord_data_g = df['cord_data']

    for chrom_str in chromosomes_list:
        if chrom_str in cord_data_g.keys():
            g = cord_data_g[chrom_str]
            betas = g['betas'][...]
            n_snps = len(betas)
            num_snps += n_snps
            sum_beta2s += sp.sum(betas ** 2)

    L = ld_scores_dict['avg_gw_ld_score']
    chi_square_lambda = sp.mean(n * sum_beta2s / float(num_snps))
    print 'Genome-wide lambda inflation:', chi_square_lambda,
    print 'Genome-wide mean LD score:', L
    gw_h2_ld_score_est = max(0.0001, (max(1, chi_square_lambda) - 1) / (n * (L / num_snps)))
    print 'Estimated genome-wide heritability:', gw_h2_ld_score_est

    assert chi_square_lambda > 1, 'Something is wrong with the GWAS summary statistics.  Perhaps there were issues parsing of them, or the given GWAS sample size (N) was too small. Either way, lambda (the mean Chi-square statistic) is too small.  '

    LDpred_inf_chrom_dict = {}
    print 'Calculating LDpred-inf weights'
    for chrom_str in chromosomes_list:
        if chrom_str in cord_data_g.keys():
            print 'Calculating scores for Chromosome %s' % ((chrom_str.split('_'))[1])
            g = cord_data_g[chrom_str]

            # Filter monomorphic SNPs
            snp_stds = g['snp_stds_ref'][...]
            snp_stds = snp_stds.flatten()
            ok_snps_filter = snp_stds > 0
            pval_derived_betas = g['betas'][...]
            n_snps = len(pval_derived_betas)
            pval_derived_betas = pval_derived_betas[ok_snps_filter]
            if h2 is not None:
                h2_chrom = h2 * (n_snps / float(num_snps))
            else:
                h2_chrom = gw_h2_ld_score_est * (n_snps / float(num_snps))
            start_betas = LDpred_inf.ldpred_inf(pval_derived_betas, genotypes=None, reference_ld_mats=chrom_ref_ld_mats[chrom_str],
                                                h2=h2_chrom, n=n, ld_window_size=2 * ld_radius, verbose=False)
            LDpred_inf_chrom_dict[chrom_str] = start_betas

    for p in ps:
        print 'Starting LDpred with p=%0.4f' % p
        p_str = '%0.4f' % p
        results_dict[p_str] = {}

        if out_file_prefix:
            # Preparing output files
            raw_effect_sizes = []
            ldpred_effect_sizes = []
            ldpred_inf_effect_sizes = []
            out_sids = []
            chromosomes = []
            out_positions = []
            out_nts = []

        for chrom_str in chromosomes_list:
            if chrom_str in cord_data_g.keys():
                g = cord_data_g[chrom_str]
                if has_phenotypes:
                    if 'raw_snps_val' in g.keys():
                        raw_snps = g['raw_snps_val'][...]
                    else:
                        raw_snps = g['raw_snps_ref'][...]

                # Filter monomorphic SNPs
                snp_stds = g['snp_stds_ref'][...]
                snp_stds = snp_stds.flatten()
                ok_snps_filter = snp_stds > 0
                snp_stds = snp_stds[ok_snps_filter]
                pval_derived_betas = g['betas'][...]
                pval_derived_betas = pval_derived_betas[ok_snps_filter]
                positions = g['positions'][...]
                positions = positions[ok_snps_filter]
                sids = g['sids'][...]
                sids = sids[ok_snps_filter]
                log_odds = g['log_odds'][...]
                log_odds = log_odds[ok_snps_filter]
                nts = g['nts'][...]
                nts = nts[ok_snps_filter]

                if out_file_prefix:
                    chromosomes.extend([chrom_str] * len(pval_derived_betas))
                    out_positions.extend(positions)
                    out_sids.extend(sids)
                    raw_effect_sizes.extend(log_odds)
                    out_nts.extend(nts)

                n_snps = len(pval_derived_betas)

                if h2 is not None:
                    h2_chrom = h2 * (n_snps / float(num_snps))
                else:
                    h2_chrom = gw_h2_ld_score_est * (n_snps / float(num_snps))
                if 'chrom_ld_boundaries' in ld_dict.keys():
                    ld_boundaries = ld_dict['chrom_ld_boundaries'][chrom_str]
                    res_dict = ldpred_gibbs(pval_derived_betas, h2=h2_chrom, n=n, p=p, ld_radius=ld_radius,
                                            verbose=verbose, num_iter=num_iter, burn_in=burn_in, ld_dict=chrom_ld_dict[chrom_str],
                                            start_betas=LDpred_inf_chrom_dict[chrom_str], ld_boundaries=ld_boundaries,
                                            zero_jump_prob=zero_jump_prob)
                else:
                    res_dict = ldpred_gibbs(pval_derived_betas, h2=h2_chrom, n=n, p=p, ld_radius=ld_radius,
                                            verbose=verbose, num_iter=num_iter, burn_in=burn_in, ld_dict=chrom_ld_dict[chrom_str],
                                            start_betas=LDpred_inf_chrom_dict[chrom_str], zero_jump_prob=zero_jump_prob)

                updated_betas = res_dict['betas']
                updated_inf_betas = res_dict['inf_betas']
                sum_sqr_effects = sp.sum(updated_betas ** 2)
                if sum_sqr_effects > gw_h2_ld_score_est:
                    print 'Sum of squared updated effects estimates seems too large:', sum_sqr_effects
                    print 'This suggests that the Gibbs sampler did not convergence.'

                print 'Calculating scores for Chromosome %s' % ((chrom_str.split('_'))[1])
                updated_betas = updated_betas / (snp_stds.flatten())
                updated_inf_betas = updated_inf_betas / (snp_stds.flatten())
                ldpred_effect_sizes.extend(updated_betas)
                ldpred_inf_effect_sizes.extend(updated_inf_betas)
                if has_phenotypes:
                    prs = sp.dot(updated_betas, raw_snps)
                    risk_scores_pval_derived += prs
                    corr = sp.corrcoef(y, prs)[0, 1]
                    r2 = corr ** 2
                    print 'The R2 prediction accuracy of PRS using %s was: %0.4f' % (chrom_str, r2)

        print 'There were %d (SNP) effects' % num_snps
        if has_phenotypes:
            num_indivs = len(y)
            results_dict[p_str]['y'] = y
            results_dict[p_str]['risk_scores_pd'] = risk_scores_pval_derived
            print 'Prediction accuracy was assessed using %d individuals.' % (num_indivs)

            corr = sp.corrcoef(y, risk_scores_pval_derived)[0, 1]
            r2 = corr ** 2
            results_dict[p_str]['r2_pd'] = r2
            print 'The  R2 prediction accuracy (observed scale) for the whole genome was: %0.4f (%0.6f)' % (r2, ((1 - r2) ** 2) / num_indivs)

            if corr < 0:
                risk_scores_pval_derived = -1 * risk_scores_pval_derived
            auc = util.calc_auc(y, risk_scores_pval_derived)
            print 'AUC for the whole genome was: %0.4f' % auc

            # Now calibration
            denominator = sp.dot(risk_scores_pval_derived.T, risk_scores_pval_derived)
            y_norm = (y - sp.mean(y)) / sp.std(y)
            numerator = sp.dot(risk_scores_pval_derived.T, y_norm)
            regression_slope = (numerator / denominator)  # [0][0]
            print 'The slope for predictions with P-value derived  effects is:', regression_slope
            results_dict[p_str]['slope_pd'] = regression_slope

        weights_out_file = '%s_LDpred_p%0.4e.txt' % (out_file_prefix, p)
        with open(weights_out_file, 'w') as f:
            f.write('chrom    pos    sid    nt1    nt2    raw_beta     ldpred_beta\n')
            for chrom, pos, sid, nt, raw_beta, ldpred_beta in it.izip(chromosomes, out_positions, out_sids, out_nts, raw_effect_sizes, ldpred_effect_sizes):
                nt1, nt2 = nt[0], nt[1]
                f.write('%s    %d    %s    %s    %s    %0.4e    %0.4e\n' % (chrom, pos, sid, nt1, nt2, raw_beta, ldpred_beta))

    weights_out_file = '%s_LDpred-inf.txt' % (out_file_prefix)
    with open(weights_out_file, 'w') as f:
        f.write('chrom    pos    sid    nt1    nt2    raw_beta    ldpred_inf_beta \n')
        for chrom, pos, sid, nt, raw_beta, ldpred_inf_beta in it.izip(chromosomes, out_positions, out_sids, out_nts, raw_effect_sizes, ldpred_inf_effect_sizes):
            nt1, nt2 = nt[0], nt[1]
            f.write('%s    %d    %s    %s    %s    %0.4e    %0.4e\n' % (chrom, pos, sid, nt1, nt2, raw_beta, ldpred_inf_beta))

def ldpred_gibbs(beta_hats, genotypes=None, start_betas=None, h2=None, n=1000, ld_radius=100,
                 num_iter=60, burn_in=10, p=None, zero_jump_prob=0.05,
                 ld_dict=None, reference_ld_mats=None, ld_boundaries=None, verbose=False):
    """
    LDpred (Gibbs Sampler)
    """
    t0 = time.time()
    m = len(beta_hats)

    # If no starting values for effects were given, then use the infinitesimal model starting values.
    if start_betas is None:
        print 'Initializing LDpred effects with posterior mean LDpred-inf effects.'
        print 'Calculating LDpred-inf effects.'
        start_betas = LDpred_inf.ldpred_inf(beta_hats, genotypes=genotypes, reference_ld_mats=reference_ld_mats,
                                            h2=h2, n=n, ld_window_size=2 * ld_radius, verbose=False)
    curr_betas = sp.copy(start_betas)
    curr_post_means = sp.zeros(m)
    avg_betas = sp.zeros(m)

    # Iterating over effect estimates in sequential order
    iter_order = sp.arange(m)

    # Setting up the marginal Bayes shrink
    Mp = m * p
    hdmp = (h2 / Mp)
    hdmpn = hdmp + 1.0 / n
    hdmp_hdmpn = (hdmp / hdmpn)
    c_const = (p / sp.sqrt(hdmpn))
    d_const = (1 - p) / (sp.sqrt(1.0 / n))

    for k in range(num_iter):  # Big iteration

        # Force an alpha shrink if estimates are way off compared to heritability estimates.  (Improves MCMC convergence.)
        h2_est = max(0.00001, sp.sum(curr_betas ** 2))
        alpha = min(1 - zero_jump_prob, 1.0 / h2_est, (h2 + 1 / sp.sqrt(n)) / h2_est)

        rand_ps = sp.random.random(m)
        rand_norms = stats.norm.rvs(0, (hdmp_hdmpn) * (1 / n), size=m)

        if ld_boundaries is None:
            for i, snp_i in enumerate(iter_order):
                start_i = max(0, snp_i - ld_radius)
                focal_i = min(ld_radius, snp_i)
                stop_i = min(m, snp_i + ld_radius + 1)

                # Local LD matrix
                D_i = ld_dict[snp_i]

                # Local (most recently updated) effect estimates
                local_betas = curr_betas[start_i: stop_i]

                # Calculate the local posterior mean, used when sampling.
                local_betas[focal_i] = 0
                res_beta_hat_i = beta_hats[snp_i] - sp.dot(D_i , local_betas)
                b2 = res_beta_hat_i ** 2

                d_const_b2_exp = d_const * sp.exp(-b2 * n / 2.0)
                if sp.isreal(d_const_b2_exp):
                    numerator = c_const * sp.exp(-b2 / (2.0 * hdmpn))
                    if sp.isreal(numerator):
                        if numerator == 0:
                            postp = 0
                        else:
                            postp = numerator / (numerator + d_const_b2_exp)
                            assert sp.isreal(postp), 'The posterior mean is not a real number?  Possibly due to problems with summary stats, LD estimates, or parameter settings.'
                    else:
                        postp = 0
                else:
                    postp = 1
                curr_post_means[snp_i] = hdmp_hdmpn * postp * res_beta_hat_i

                if rand_ps[i] < postp * alpha:
                    # Sample from the posterior Gaussian dist.
                    proposed_beta = rand_norms[i] + hdmp_hdmpn * res_beta_hat_i

                else:
                    # Sample 0
                    proposed_beta = 0

                curr_betas[snp_i] = proposed_beta  # UPDATE BETA
        else:
            for i, snp_i in enumerate(iter_order):
                start_i = ld_boundaries[snp_i][0]
                stop_i = ld_boundaries[snp_i][1]
                focal_i = snp_i - start_i

                # Local LD matrix
                D_i = ld_dict[snp_i]

                # Local (most recently updated) effect estimates
                local_betas = curr_betas[start_i: stop_i]

                # Calculate the local posterior mean, used when sampling.
                local_betas[focal_i] = 0
                res_beta_hat_i = beta_hats[snp_i] - sp.dot(D_i , local_betas)
                b2 = res_beta_hat_i ** 2

                d_const_b2_exp = d_const * sp.exp(-b2 * n / 2.0)
                if sp.isreal(d_const_b2_exp):
                    numerator = c_const * sp.exp(-b2 / (2.0 * hdmpn))
                    if sp.isreal(numerator):
                        if numerator == 0:
                            postp = 0
                        else:
                            postp = numerator / (numerator + d_const_b2_exp)
                            assert sp.isreal(postp), 'Posterior mean is not a real number? Possibly due to problems with summary stats, LD estimates, or parameter settings.'
                    else:
                        postp = 0
                else:
                    postp = 1
                curr_post_means[snp_i] = hdmp_hdmpn * postp * res_beta_hat_i

                if rand_ps[i] < postp * alpha:
                    # Sample from the posterior Gaussian dist.
                    proposed_beta = rand_norms[i] + hdmp_hdmpn * res_beta_hat_i

                else:
                    # Sample 0
                    proposed_beta = 0

                curr_betas[snp_i] = proposed_beta  # UPDATE BETA
        if verbose:
            sys.stdout.write('\b\b\b\b\b\b\b%0.2f%%' % (100.0 * (min(1, float(k + 1) / num_iter))))
            sys.stdout.flush()

        if k >= burn_in:
            avg_betas += curr_post_means  # Averaging over the posterior means instead of samples.

    avg_betas = avg_betas / float(num_iter - burn_in)
    t1 = time.time()
    t = (t1 - t0)
    if verbose:
        print '\nTook %d minutes and %0.2f seconds' % (t / 60, t % 60)
    return {'betas':avg_betas, 'inf_betas':start_betas}

def main():
    p_dict = parse_parameters()
    local_ld_dict_file = '%s_ldradius%d.pickled.gz' % (p_dict['local_ld_file_prefix'], p_dict['ld_radius'])

    print """
Note: For maximal accuracy all SNPs with LDpred weights should be included in the validation data set.
If they are a subset of the validation data set, then we suggest recalculate LDpred for the overlapping SNPs.
"""
    if not os.path.isfile(local_ld_dict_file):
        df = h5py.File(p_dict['coord'])

        chrom_ld_scores_dict = {}
        chrom_ld_dict = {}
        chrom_ref_ld_mats = {}
        if p_dict['gm_ld_radius'] is not None:
            chrom_ld_boundaries = {}
        ld_score_sum = 0
        num_snps = 0
        print 'Calculating LD information w. radius %d' % p_dict['ld_radius']

        cord_data_g = df['cord_data']

        for chrom_str in cord_data_g.keys():
            print 'Working on %s' % chrom_str
            g = cord_data_g[chrom_str]
            if 'raw_snps_ref' in g.keys():
                raw_snps = g['raw_snps_ref'][...]
                snp_stds = g['snp_stds_ref'][...]
                snp_means = g['snp_means_ref'][...]

            # Filter monomorphic SNPs
            ok_snps_filter = snp_stds > 0
            ok_snps_filter = ok_snps_filter.flatten()
            raw_snps = raw_snps[ok_snps_filter]
            snp_means = snp_means[ok_snps_filter]
            snp_stds = snp_stds[ok_snps_filter]

            n_snps = len(raw_snps)
            snp_means.shape = (n_snps, 1)
            snp_stds.shape = (n_snps, 1)

            # Normalize SNPs..
            snps = sp.array((raw_snps - snp_means) / snp_stds, dtype='float32')
            assert snps.shape == raw_snps.shape, 'Problems normalizing SNPs (array shape mismatch).'

            # ??? Other stuff???
            if p_dict['gm_ld_radius'] is not None:
                assert 'genetic_map' in g.keys(), 'Genetic map is missing.'
                gm = g['genetic_map'][...]
                ret_dict = ld.get_LDpred_ld_tables(snps, gm=gm, gm_ld_radius=p_dict['gm_ld_radius'])
                chrom_ld_boundaries[chrom_str] = ret_dict['ld_boundaries']
            else:
                ret_dict = ld.get_LDpred_ld_tables(snps, ld_radius=p_dict['ld_radius'], ld_window_size=2 * p_dict['ld_radius'])
            chrom_ld_dict[chrom_str] = ret_dict['ld_dict']
            chrom_ref_ld_mats[chrom_str] = ret_dict['ref_ld_matrices']
            ld_scores = ret_dict['ld_scores']
            chrom_ld_scores_dict[chrom_str] = {'ld_scores':ld_scores, 'avg_ld_score':sp.mean(ld_scores)}
            ld_score_sum += sp.sum(ld_scores)
            num_snps += n_snps
        avg_gw_ld_score = ld_score_sum / float(num_snps)
        ld_scores_dict = {'avg_gw_ld_score': avg_gw_ld_score, 'chrom_dict':chrom_ld_scores_dict}

        print 'Done calculating the LD table and LD score, writing to file:', local_ld_dict_file
        print 'Genome-wide average LD score was:', ld_scores_dict['avg_gw_ld_score']
        ld_dict = {'ld_scores_dict':ld_scores_dict, 'chrom_ld_dict':chrom_ld_dict, 'chrom_ref_ld_mats':chrom_ref_ld_mats}
        if p_dict['gm_ld_radius'] is not None:
            ld_dict['chrom_ld_boundaries'] = chrom_ld_boundaries
        f = gzip.open(local_ld_dict_file, 'wb')
        cPickle.dump(ld_dict, f, protocol=2)
        f.close()
        print 'LD information is now pickled.'
    else:
        print 'Loading LD information from file: %s' % local_ld_dict_file
        f = gzip.open(local_ld_dict_file, 'r')
        ld_dict = cPickle.load(f)
        f.close()

    ldpred_genomewide(data_file=p_dict['coord'], out_file_prefix=p_dict['out'], ps=p_dict['PS'], ld_radius=p_dict['ld_radius'],
                      ld_dict=ld_dict, n=p_dict['N'], num_iter=p_dict['num_iter'], h2=p_dict['H2'], verbose=False)

if __name__ == '__main__':
    main()


#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Restricted Open-shell Hartree-Fock
'''

from functools import reduce
import numpy
import pyscf.gto
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import hf
from pyscf.scf import uhf
import pyscf.scf.chkfile
from pyscf import __config__

WITH_META_LOWDIN = getattr(__config__, 'scf_analyze_with_meta_lowdin', True)
MO_BASE = getattr(__config__, 'MO_BASE', 1)

def init_guess_by_minao(mol):
    dm = hf.init_guess_by_minao(mol)
    return numpy.array((dm*.5, dm*.5))

def init_guess_by_atom(mol):
    dm = hf.init_guess_by_atom(mol)
    return numpy.array((dm*.5, dm*.5))

init_guess_by_huckel = hf.init_guess_by_huckel
init_guess_by_chkfile = uhf.init_guess_by_chkfile

def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
             diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
    '''Build fock matrix based on Roothaan's effective fock.
    See also :func:`get_roothaan_fock`
    '''
    if h1e is None: h1e = mf.get_hcore()
    if s1e is None: s1e = mf.get_ovlp()
    if vhf is None: vhf = mf.get_veff(mf.mol, dm)
    if dm is None: dm = mf.make_rdm1()
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        dm = numpy.array((dm*.5, dm*.5))
# To Get orbital energy in get_occ, we saved alpha and beta fock, because
# Roothaan effective Fock cannot provide correct orbital energy with `eig`
# TODO, check other treatment  J. Chem. Phys. 133, 141102
    focka = h1e + vhf[0]
    fockb = h1e + vhf[1]
    f = get_roothaan_fock((focka,fockb), dm, s1e)
    if cycle < 0 and diis is None:  # Not inside the SCF iteration
        return f

    if diis_start_cycle is None:
        diis_start_cycle = mf.diis_start_cycle
    if level_shift_factor is None:
        level_shift_factor = mf.level_shift
    if damp_factor is None:
        damp_factor = mf.damp

    dm_tot = dm[0] + dm[1]
    if 0 <= cycle < diis_start_cycle-1 and abs(damp_factor) > 1e-4:
        raise NotImplementedError('ROHF Fock-damping')
    if diis and cycle >= diis_start_cycle:
        f = diis.update(s1e, dm_tot, f, mf, h1e, vhf)
    if abs(level_shift_factor) > 1e-4:
        f = hf.level_shift(s1e, dm_tot*.5, f, level_shift_factor)
    f = lib.tag_array(f, focka=focka, fockb=fockb)
    return f

def get_roothaan_fock(focka_fockb, dma_dmb, s):
    '''Roothaan's effective fock.
    Ref. http://www-theor.ch.cam.ac.uk/people/ross/thesis/node15.html

    ======== ======== ====== =========
    space     closed   open   virtual
    ======== ======== ====== =========
    closed      Fc      Fb     Fc
    open        Fb      Fc     Fa
    virtual     Fc      Fa     Fc
    ======== ======== ====== =========

    where Fc = (Fa + Fb) / 2

    Returns:
        Roothaan effective Fock matrix
    '''
    nao = s.shape[0]
    focka, fockb = focka_fockb
    dma, dmb = dma_dmb
    fc = (focka + fockb) * .5
# Projector for core, open-shell, and virtual
    pc = numpy.dot(dmb, s)
    po = numpy.dot(dma-dmb, s)
    pv = numpy.eye(nao) - numpy.dot(dma, s)
    fock  = reduce(numpy.dot, (pc.conj().T, fc, pc)) * .5
    fock += reduce(numpy.dot, (po.conj().T, fc, po)) * .5
    fock += reduce(numpy.dot, (pv.conj().T, fc, pv)) * .5
    fock += reduce(numpy.dot, (po.conj().T, fockb, pc))
    fock += reduce(numpy.dot, (po.conj().T, focka, pv))
    fock += reduce(numpy.dot, (pv.conj().T, fc, pc))
    fock = fock + fock.conj().T
    fock = lib.tag_array(fock, focka=focka, fockb=fockb)
    return fock

def get_occ(mf, mo_energy=None, mo_coeff=None):
    '''Label the occupancies for each orbital.
    NOTE the occupancies are not assigned based on the orbital energy ordering.
    The first N orbitals are assigned to be occupied orbitals.

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; O 0 0 1.1', spin=1)
    >>> mf = scf.hf.SCF(mol)
    >>> energy = numpy.array([-10., -1., 1, -2., 0, -3])
    >>> mf.get_occ(energy)
    array([2, 2, 2, 2, 1, 0])
    '''

    if mo_energy is None: mo_energy = mf.mo_energy
    if getattr(mo_energy, 'mo_ea', None) is not None:
        mo_ea = mo_energy.mo_ea
        mo_eb = mo_energy.mo_eb
    else:
        mo_ea = mo_eb = mo_energy
    nmo = mo_ea.size
    mo_occ = numpy.zeros(nmo)
    if getattr(mf, 'nelec', None) is None:
        nelec = mf.mol.nelec
    else:
        nelec = mf.nelec
    ncore = nelec[1]
    nocc  = nelec[0]
    nopen = abs(nocc - ncore)
    mo_occ = _fill_rohf_occ(mo_energy, mo_ea, mo_eb, ncore, nopen)

    if mf.verbose >= logger.INFO and nocc < nmo and ncore > 0:
        ehomo = max(mo_energy[mo_occ> 0])
        elumo = min(mo_energy[mo_occ==0])
        if ehomo+1e-3 > elumo:
            logger.warn(mf, 'HOMO %.15g >= LUMO %.15g', ehomo, elumo)
        else:
            logger.info(mf, '  HOMO = %.15g  LUMO = %.15g', ehomo, elumo)
        if nopen > 0 and mf.verbose >= logger.DEBUG:
            core_idx = mo_occ == 2
            open_idx = mo_occ == 1
            vir_idx = mo_occ == 0
            logger.debug(mf, '                  Roothaan           | alpha              | beta')
            logger.debug(mf, '  Highest 2-occ = %18.15g | %18.15g | %18.15g',
                         max(mo_energy[core_idx]),
                         max(mo_ea[core_idx]), max(mo_eb[core_idx]))
            logger.debug(mf, '  Lowest 0-occ =  %18.15g | %18.15g | %18.15g',
                         min(mo_energy[vir_idx]),
                         min(mo_ea[vir_idx]), min(mo_eb[vir_idx]))
            for i in numpy.where(open_idx)[0]:
                logger.debug(mf, '  1-occ =         %18.15g | %18.15g | %18.15g',
                             mo_energy[i], mo_ea[i], mo_eb[i])

        if mf.verbose >= logger.DEBUG:
            numpy.set_printoptions(threshold=nmo)
            logger.debug(mf, '  Roothaan mo_energy =\n%s', mo_energy)
            logger.debug1(mf, '  alpha mo_energy =\n%s', mo_ea)
            logger.debug1(mf, '  beta  mo_energy =\n%s', mo_eb)
            numpy.set_printoptions(threshold=1000)
    return mo_occ

def _fill_rohf_occ(mo_energy, mo_energy_a, mo_energy_b, ncore, nopen):
    mo_occ = numpy.zeros_like(mo_energy)
    open_idx = []
    core_sort = numpy.argsort(mo_energy)
    core_idx = core_sort[:ncore]
    if nopen > 0:
        open_idx = core_sort[ncore:]
        open_sort = numpy.argsort(mo_energy_a[open_idx])
        open_idx = open_idx[open_sort[:nopen]]
    mo_occ[core_idx] = 2
    mo_occ[open_idx] = 1
    return mo_occ

def get_grad(mo_coeff, mo_occ, fock):
    '''ROHF gradients is the off-diagonal block [co + cv + ov], where
    [ cc co cv ]
    [ oc oo ov ]
    [ vc vo vv ]
    '''
    occidxa = mo_occ > 0
    occidxb = mo_occ == 2
    viridxa = ~occidxa
    viridxb = ~occidxb
    uniq_var_a = viridxa.reshape(-1,1) & occidxa
    uniq_var_b = viridxb.reshape(-1,1) & occidxb

    if getattr(fock, 'focka', None) is not None:
        focka = fock.focka
        fockb = fock.fockb
    elif isinstance(fock, (tuple, list)) or getattr(fock, 'ndim', None) == 3:
        focka, fockb = fock
    else:
        focka = fockb = fock
    focka = reduce(numpy.dot, (mo_coeff.conj().T, focka, mo_coeff))
    fockb = reduce(numpy.dot, (mo_coeff.conj().T, fockb, mo_coeff))

    g = numpy.zeros_like(focka)
    g[uniq_var_a]  = focka[uniq_var_a]
    g[uniq_var_b] += fockb[uniq_var_b]
    return g[uniq_var_a | uniq_var_b]

def make_rdm1(mo_coeff, mo_occ, **kwargs):
    '''One-particle densit matrix.  mo_occ is a 1D array, with occupancy 1 or 2.
    '''
    mo_a = mo_coeff[:,mo_occ>0]
    mo_b = mo_coeff[:,mo_occ==2]
    dm_a = numpy.dot(mo_a, mo_a.conj().T)
    dm_b = numpy.dot(mo_b, mo_b.conj().T)
    return numpy.array((dm_a, dm_b))

def energy_elec(mf, dm=None, h1e=None, vhf=None):
    if dm is None: dm = mf.make_rdm1()
    elif isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        dm = numpy.array((dm*.5, dm*.5))
    return uhf.energy_elec(mf, dm, h1e, vhf)

get_veff = uhf.get_veff

def analyze(mf, verbose=logger.DEBUG, with_meta_lowdin=WITH_META_LOWDIN,
            **kwargs):
    '''Analyze the given SCF object:  print orbital energies, occupancies;
    print orbital coefficients; Mulliken population analysis
    '''
    from pyscf.lo import orth
    from pyscf.tools import dump_mat
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    mo_coeff = mf.mo_coeff
    log = logger.new_logger(mf, verbose)
    if log.verbose >= logger.NOTE:
        log.note('**** MO energy ****')
        if getattr(mo_energy, 'mo_ea', None) is not None:
            mo_ea = mo_energy.mo_ea
            mo_eb = mo_energy.mo_eb
            log.note('                Roothaan           | alpha              | beta')
            for i,c in enumerate(mo_occ):
                log.note('MO #%-3d energy= %-18.15g | %-18.15g | %-18.15g occ= %g',
                         i+MO_BASE, mo_energy[i], mo_ea[i], mo_eb[i], c)
        else:
            for i,c in enumerate(mo_occ):
                log.note('MO #%-3d energy= %-18.15g occ= %g',
                         i+MO_BASE, mo_energy[i], c)

    ovlp_ao = mf.get_ovlp()
    if log.verbose >= logger.DEBUG:
        label = mf.mol.ao_labels()
        if with_meta_lowdin:
            log.debug(' ** MO coefficients (expansion on meta-Lowdin AOs) **')
            orth_coeff = orth.orth_ao(mf.mol, 'meta_lowdin', s=ovlp_ao)
            c = reduce(numpy.dot, (orth_coeff.conj().T, ovlp_ao, mo_coeff))
        else:
            log.debug(' ** MO coefficients (expansion on AOs) **')
            c = mo_coeff
        dump_mat.dump_rec(mf.stdout, c, label, start=MO_BASE, **kwargs)
    dm = mf.make_rdm1(mo_coeff, mo_occ)
    if with_meta_lowdin:
        pop_and_charge = mf.mulliken_meta(mf.mol, dm, s=ovlp_ao, verbose=log)
    else:
        pop_and_charge = mf.mulliken_pop(mf.mol, dm, s=ovlp_ao, verbose=log)
    dip = mf.dip_moment(mf.mol, dm, verbose=log)
    return pop_and_charge, dip

mulliken_pop = hf.mulliken_pop
mulliken_meta = hf.mulliken_meta

def canonicalize(mf, mo_coeff, mo_occ, fock=None):
    '''Canonicalization diagonalizes the Fock matrix within occupied, open,
    virtual subspaces separatedly (without change occupancy).
    '''
    if getattr(fock, 'focka', None) is None:
        dm = mf.make_rdm1(mo_coeff, mo_occ)
        fock = mf.get_fock(dm=dm)
    mo_e, mo_coeff = hf.canonicalize(mf, mo_coeff, mo_occ, fock)
    fa, fb = fock.focka, fock.fockb
    mo_ea = numpy.einsum('pi,pi->i', mo_coeff.conj(), fa.dot(mo_coeff)).real
    mo_eb = numpy.einsum('pi,pi->i', mo_coeff.conj(), fb.dot(mo_coeff)).real
    mo_e = lib.tag_array(mo_e, mo_ea=mo_ea, mo_eb=mo_eb)
    return mo_e, mo_coeff

dip_moment = hf.dip_moment

# use UHF init_guess, get_veff, diis, and intermediates such as fock, vhf, dm
# keep mo_energy, mo_coeff, mo_occ as RHF structure

class ROHF(hf.RHF):
    __doc__ = hf.SCF.__doc__

    def __init__(self, mol):
        hf.SCF.__init__(self, mol)
        self.nelec = None

    @property
    def nelec(self):
        if getattr(self, '_nelec', None) is not None:
            return self._nelec
        else:
            return self.mol.nelec
    @nelec.setter
    def nelec(self, x):
        self._nelec = x

    @property
    def nelectron_alpha(self):
        return self.nelec[0]
    @nelectron_alpha.setter
    def nelectron_alpha(self, x):
        logger.warn(self, 'WARN: Attribute .nelectron_alpha is deprecated. '
                    'Set .nelec instead')
        #raise RuntimeError('API updates')
        self.nelec = (x, self.mol.nelectron-x)

    check_sanity = hf.SCF.check_sanity

    def dump_flags(self, verbose=None):
        hf.SCF.dump_flags(self, verbose)
        nelec = self.nelec
        logger.info(self, 'num. doubly occ = %d  num. singly occ = %d',
                    nelec[1], nelec[0]-nelec[1])

    def init_guess_by_minao(self, mol=None):
        if mol is None: mol = self.mol
        return init_guess_by_minao(mol)

    def init_guess_by_atom(self, mol=None):
        if mol is None: mol = self.mol
        logger.info(self, 'Initial guess from the superpostion of atomic densties.')
        return init_guess_by_atom(mol)

    def init_guess_by_huckel(self, mol=None):
        if mol is None: mol = self.mol
        logger.info(self, 'Initial guess from on-the-fly Huckel, doi:10.1021/acs.jctc.8b01089.')
        mo_energy, mo_coeff = init_guess_by_huckel(mol)
        mo_occ = self.get_occ(mo_energy, mo_coeff)
        return self.make_rdm1(mo_coeff, mo_occ)

    def init_guess_by_1e(self, mol=None):
        if mol is None: mol = self.mol
        logger.info(self, 'Initial guess from hcore.')
        h1e = self.get_hcore(mol)
        s1e = self.get_ovlp(mol)
        mo_energy, mo_coeff = self.eig(h1e, s1e)
        mo_occ = self.get_occ(mo_energy, mo_coeff)
        return self.make_rdm1(mo_coeff, mo_occ)

    def init_guess_by_chkfile(self, chkfile=None, project=None):
        if chkfile is None: chkfile = self.chkfile
        return init_guess_by_chkfile(self.mol, chkfile, project=project)

    get_fock = get_fock
    get_occ = get_occ

    @lib.with_doc(hf.eig.__doc__)
    def eig(self, fock, s):
        e, c = self._eigh(fock, s)
        if getattr(fock, 'focka', None) is not None:
            mo_ea = numpy.einsum('pi,pi->i', c.conj(), fock.focka.dot(c)).real
            mo_eb = numpy.einsum('pi,pi->i', c.conj(), fock.fockb.dot(c)).real
            e = lib.tag_array(e, mo_ea=mo_ea, mo_eb=mo_eb)
        return e, c

    @lib.with_doc(get_grad.__doc__)
    def get_grad(self, mo_coeff, mo_occ, fock=None):
        if fock is None:
            dm1 = self.make_rdm1(mo_coeff, mo_occ)
            fock = self.get_hcore(self.mol) + self.get_veff(self.mol, dm1)
        return get_grad(mo_coeff, mo_occ, fock)

    @lib.with_doc(make_rdm1.__doc__)
    def make_rdm1(self, mo_coeff=None, mo_occ=None, **kwargs):
        if mo_coeff is None: mo_coeff = self.mo_coeff
        if mo_occ is None: mo_occ = self.mo_occ
        return make_rdm1(mo_coeff, mo_occ, **kwargs)

    energy_elec = energy_elec

    @lib.with_doc(uhf.get_veff.__doc__)
    def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
            dm = numpy.array((dm*.5, dm*.5))

        if self._eri is not None or not self.direct_scf:
            if getattr(dm, 'mo_coeff', None) is not None:
                mo_coeff = dm.mo_coeff
                mo_occ_a = (dm.mo_occ > 0).astype(numpy.double)
                mo_occ_b = (dm.mo_occ ==2).astype(numpy.double)
                dm = lib.tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
                                   mo_occ=(mo_occ_a,mo_occ_b))
            vj, vk = self.get_jk(mol, dm, hermi)
            vhf = vj[0] + vj[1] - vk
        else:
            ddm = dm - numpy.asarray(dm_last)
            vj, vk = self.get_jk(mol, ddm, hermi)
            vhf = vj[0] + vj[1] - vk
            vhf += numpy.asarray(vhf_last)
        return vhf

    @lib.with_doc(analyze.__doc__)
    def analyze(self, verbose=None, with_meta_lowdin=WITH_META_LOWDIN,
                **kwargs):
        if verbose is None: verbose = self.verbose
        return analyze(self, verbose, with_meta_lowdin, **kwargs)

    canonicalize = canonicalize

    def spin_square(self, mo_coeff=None, s=None):
        '''Spin square and multiplicity of RHF determinant'''
        neleca, nelecb = self.nelec
        ms = (neleca - nelecb) * .5
        ss = ms * (ms + 1)
        return ss, ms*2+1

    def stability(self,
                  internal=getattr(__config__, 'scf_stability_internal', True),
                  external=getattr(__config__, 'scf_stability_external', False),
                  verbose=None):
        '''
        ROHF/ROKS stability analysis.

        See also pyscf.scf.stability.rohf_stability function.

        Kwargs:
            internal : bool
                Internal stability, within the RHF optimization space.
            external : bool
                External stability. It is not available in current version.

        Returns:
            The return value includes two set of orbitals which are more close to
            the required stable condition.
        '''
        from pyscf.scf.stability import rohf_stability
        return rohf_stability(self, internal, external, verbose)

    def nuc_grad_method(self):
        from pyscf.grad import rohf
        return rohf.Gradients(self)

class HF1e(ROHF):
    def scf(self, *args):
        logger.info(self, '\n')
        logger.info(self, '******** 1 electron system ********')
        self.converged = True
        h1e = self.get_hcore(self.mol)
        s1e = self.get_ovlp(self.mol)
        self.mo_energy, self.mo_coeff = self.eig(h1e, s1e)
        self.mo_occ = self.get_occ(self.mo_energy, self.mo_coeff)
        self.e_tot = self.mo_energy[self.mo_occ>0][0] + self.mol.energy_nuc()
        self._finalize()
        return self.e_tot

del(WITH_META_LOWDIN)


"""
Defines the KeplerFFI class that uses FFIs to model the PRF shape for a given
channel and quarter.
"""
import os
import sys
import warnings
import wget

import numpy as np
import pandas as pd
import pickle
import matplotlib.pyplot as plt
import matplotlib.colors as colors

from scipy import sparse
from tqdm.auto import tqdm
from astropy.io import fits
from astropy.table import Table
from astropy.coordinates import SkyCoord, match_coordinates_3d
from astropy.stats import sigma_clip, SigmaClip
from astropy.time import Time
from astropy.wcs import WCS
from photutils import Background2D, MedianBackground, BkgZoomInterpolator

from . import PACKAGEDIR, DATAOUTDIR
from .utils import get_gaia_sources, make_A_edges, solve_linear_model, _make_A_polar

r_min, r_max = 20, 1044
c_min, c_max = 12, 1112
remove_sat = True
mask_bright = True

# dictionary with FFI file names and in-quarter mapping
quarter_ffi = {
    0: [
        "kplr2009114174833_ffi-cal.fits",
        "kplr2009114204835_ffi-cal.fits",
        "kplr2009115002613_ffi-cal.fits",
        "kplr2009115053616_ffi-cal.fits",
        "kplr2009115080620_ffi-cal.fits",
        "kplr2009115131122_ffi-cal.fits",
        "kplr2009115173611_ffi-cal.fits",
        "kplr2009116035924_ffi-cal.fits",
    ],
    1: [],
    2: ["kplr2009231194831_ffi-cal.fits"],
    3: ["kplr2009292020429_ffi-cal.fits", "kplr2009322233047_ffi-cal.fits"],
    4: [
        "kplr2010019225502_ffi-cal.fits",
        "kplr2010020005046_ffi-cal.fits",
        "kplr2010049182302_ffi-cal.fits",
    ],
    5: ["kplr2010111125026_ffi-cal.fits", "kplr2010140101631_ffi-cal.fits"],
    6: ["kplr2010203012215_ffi-cal.fits", "kplr2010234192745_ffi-cal.fits"],
    7: ["kplr2010296192119_ffi-cal.fits", "kplr2010326181728_ffi-cal.fits"],
    8: ["kplr2011024134926_ffi-cal.fits", "kplr2011053174401_ffi-cal.fits"],
    9: ["kplr2011116104002_ffi-cal.fits", "kplr2011145152723_ffi-cal.fits"],
    10: ["kplr2011208112727_ffi-cal.fits", "kplr2011240181752_ffi-cal.fits"],
    11: ["kplr2011303191211_ffi-cal.fits", "kplr2011334181008_ffi-cal.fits"],
    12: ["kplr2012032101442_ffi-cal.fits", "kplr2012060123308_ffi-cal.fits"],
    13: ["kplr2012121122500_ffi-cal.fits", "kplr2012151105138_ffi-cal.fits"],
    14: ["kplr2012211123923_ffi-cal.fits", "kplr2012242195726_ffi-cal.fits"],
    15: ["kplr2012310200152_ffi-cal.fits", "kplr2012341215621_ffi-cal.fits"],
    16: ["kplr2013038133130_ffi-cal.fits", "kplr2013065115251_ffi-cal.fits"],
    17: ["kplr2013098115308_ffi-cal.fits"],
}

class KeplerFFI(object):
    """
    Class for loading Kepler's FFI files and compute PRF models out of them following
    the method discussed in Hedges et al. 2021 and Martinez-Palomera et al. 2021.
    """

    def __init__(
        self,
        ffi_name: str = "",
        channel: int = 1,
        quarter: int = None,
        plot: bool = True,
        save: bool = True,
    ):
        """
        Initialization of the KeplerFFI class
        Parameters
        ----------
        ffi_name : string
            Name of the FFI file used to model the PRF profile. Either ffi_name or
            quarter can be provided.
        channel : int
            Channel number of the FFI to be used to model the PRF. Valid values are
            between 1 and 84.
        quarter : int
            Number of the quarter that will be used to model the PRF.
            Either ffi_name or quarter can be provided, if quarter is provided,
            then all FFI files observed during the time window of that quarter will
            be used to model the PRF by averaging the images.
            Valid values are between 1 and 17.
        plot : boolean
            Whether to clreate diagnostic plots or not.
        save : boolean
            Whether to save the models or not.

        Attributes
        ----------
        channel : int
            Number of the quarter that will be used to model the PRF.
        quarter : int
            Channel number of the FFI to be used to model the PRF.
        plot : bool
            Boolean to create diagnostic plots.
        save : bool
            Boolean to save models and figures.
        hdr : dict
            Header dictionary of the FFI file.
        img : numpy.ndarray
            Original FFI flux image in electros / sec.
        wcs : astropy.WCS
            Object with the WCS solution of the image.
        col_2d : numpy.ndarray
            Data array with the pixel column number in 2D
        row_2d : numpy.ndarray
            Data array with the pixel row number in 2D
        ra_2d : numpy.ndarray
            Data array with the pixel Right Ascension value in 2D, in degs
        dec_2d : numpy.ndarray
            Data array with the pixel Declination value in 2D, in degs
        flux_2d : numpy.ndarray
            Data array with the flux value of each pixel with substracted background
            in 2D, in electros / sec.
        sources : pandas.DataFrame
            Catalog with Gaia sources observed in the image, after cleaning.
        flux : numpy.ndarray
            Data array with the flux value of each pixel with substracted background
            in 1D after removing saturated & bright pixels. in electros / sec.
        flux_err : numpy.ndarray
            Data array with the flux error value of each pixel with substracted
            background in 1D after removing saturated & bright pixels. in electros / sec.
        col : numpy.ndarray
            Data array with the pixel column number in 1D after removing saturated &
            bright pixels
        row : numpy.ndarray
            Data array with the pixel row number in 1D after removing saturated &
            bright pixels
        nsurces : int
            Total number of sources observed in the image after cleaning.
        npixels : int
            Total number of pixels in the image
        gf : numpy.ndarray
            Data array with the Gaia flux value for every source.
        dflux : scipy.sparse.csr_matrix
            Sparse matrix with pixel flux value within r < 7 pixels of the source
            coordinates. Has shape [nsources , npixels]
        dx : scipy.sparse.csr_matrix
            Sparse matrix with distance between the pixel within r < 7 pixels and the
            source location, in pixel units. Has shape [nsources , npixels]
        dy : scipy.sparse.csr_matrix
            Sparse matrix with distance between the pixel within r < 7 pixels and the
            source location, in pixel units. Has shape [nsources , npixels]
        r : scipy.sparse.csr_matrix
            Sparse matrix with radial distance between pixels within r < 7 and the
            source location, in polar coordinates. Has shape [nsources , npixels]
        phi : scipy.sparse.csr_matrix
            Sparse matrix with angle value of pixels within r < 7 and the
            source location, in polar coordinates. Has shape [nsources , npixels]
        """

        self.channel = channel
        self.plot = plot
        self.save = save
        self.show = False

        if quarter is not None and quarter in np.arange(18):
            fname = "%s/data/fits/ffi/%s" % (DATAOUTDIR, quarter_ffi[quarter][0])
        elif len(ffi_name) == 17 and ffi_name[:4] == "kplr":
            fname = "%s/data/fits/ffi/%s_ffi-cal.fits" % (DATAOUTDIR, ffi_name)
        else:
            raise ValueError("Invalid quarter or FFI fits file name")

        if not os.path.isfile(fname):
            print("Downloading FFI fits files")
            fits_name = fname.split("/")[-1]
            print(fits_name)
            self.download_ffi(fits_name)

        self.hdr = fits.open(fname)[channel].header
        self.img = fits.open(fname)[channel].data
        self.wcs = WCS(self.hdr)
        self.quarter = quarter if quarter is not None else self.hdr["MJDSTART"]

        row_2d, col_2d = np.mgrid[: self.img.shape[0], : self.img.shape[1]]
        row, col = row_2d.ravel(), col_2d.ravel()
        ra, dec = self.wcs.all_pix2world(np.vstack([col, row]).T, 0).T
        ra_2d, dec_2d = ra.reshape(self.img.shape), dec.reshape(self.img.shape)

        # get coordinates of the center for query
        loc = (self.img.shape[0] // 2, self.img.shape[1] // 2)
        ra_q, dec_q = self.wcs.all_pix2world(np.atleast_2d(loc), 0).T
        rad = [np.hypot(ra - ra.mean(), dec - dec.mean()).max()]

        time = Time(self.hdr["TSTART"] + 2454833, format="jd")
        if ra_q[0] > 360 or np.abs(dec_q[0]) > 90 or rad[0] > 5:
            raise ValueError(
                "Query values are out of bound, please check WCS solution."
            )

        # remove border Pixels
        self.col_2d = col_2d[r_min:r_max, c_min:c_max] - c_min
        self.row_2d = row_2d[r_min:r_max, c_min:c_max] - r_min
        self.ra_2d = ra_2d[r_min:r_max, c_min:c_max]
        self.dec_2d = dec_2d[r_min:r_max, c_min:c_max]
        flux_2d = self.img[r_min:r_max, c_min:c_max]

        sources = self._do_big_query(self.ra_2d, self.dec_2d, time.jyear)
        sources["col"], sources["row"] = self.wcs.all_world2pix(
            sources.loc[:, ["ra", "dec"]].values, 0.5
        ).T

        # correct col,row columns for gaia sources
        sources.row -= r_min
        sources.col -= c_min

        # clean out-of-ccd and blended sources
        clean_sources = self._clean_source_list(sources)
        del sources

        # background substraction
        self.flux_2d = flux_2d - self._model_bkg(flux_2d, mask=None)

        # ravel arrays
        col = self.col_2d.ravel()
        row = self.row_2d.ravel()
        ra = self.ra_2d.ravel()
        dec = self.dec_2d.ravel()
        flux = self.flux_2d.ravel()

        if remove_sat:
            non_sat_mask = ~self._saturated_pixels_mask(
                flux, col, row, saturation_limit=1.5e5
            )
            print("Saturated pixels %i: " % (np.sum(~non_sat_mask)))
            self.non_sat_mask = non_sat_mask

            col = col[non_sat_mask]
            row = row[non_sat_mask]
            ra = ra[non_sat_mask]
            dec = dec[non_sat_mask]
            flux = flux[non_sat_mask]

        if mask_bright:
            bright_mask = ~self._mask_bright_sources(
                flux, col, row, clean_sources, mag_limit=10
            )
            print("Bright pixels %i: " % (np.sum(~bright_mask)))
            self.bright_mask = bright_mask

            col = col[bright_mask]
            row = row[bright_mask]
            ra = ra[bright_mask]
            dec = dec[bright_mask]
            flux = flux[bright_mask]

        clean_sources = clean_sources[
            (clean_sources.phot_g_mean_flux > 1e3)
            & (clean_sources.phot_g_mean_flux < 1e6)
        ].reset_index(drop=True)

        print("Total Gaia sources %i: " % (clean_sources.shape[0]))

        self.sources = clean_sources
        self.flux = flux
        self.flux_err = np.sqrt(np.abs(self.flux))
        self.col = col
        self.row = row
        self.nsurces = clean_sources.shape[0]
        self.npixels = self.flux.shape[0]

        self.rmin = 0.25
        self.rmax = 3.0

    @staticmethod
    def download_ffi(fits_name):
        """
        Download FFI fits file to a dedicated quarter directory

        Parameters
        ----------
        fits_name : string
            Name of FFI fits file
        """
        url = "https://archive.stsci.edu/missions/kepler/ffi"
        if fits_name == "":
            raise ValueError("Invalid fits file name")

        if not os.path.isdir("%s/data/fits/ffi" % (DATAOUTDIR)):
            os.makedirs("%s/data/fits/ffi" % (DATAOUTDIR))

        out = "%s/data/fits/ffi/%s" % (DATAOUTDIR, fits_name)
        wget.download("%s/%s" % (url, fits_name), out=out)

        return

    def _do_big_query(self, ra, dec, epoch):
        """
        Query Gaia catalogs (EDR3 default) to obtain sources observed in the FFI.
        If query finishs ok, result will be saved for future use in the following
        directory:
            ../data/catalogs/ffi/<quarter#>/channel_<channel#>_gaia_xmatch.csv

        It does nx*ny small queries to avoid TimeoutError that might happen when doing
        large (rad > 0.7 deg) queries to Gaia archive. The ouput file has unique
        objects.

        Parameters
        ----------
        ra : list
            Value of the Right Ascension coordinate used for the query, in deg.
        dec : list
            Value of the Declination coordinate used for the query, in deg.
        epoch : float
            Year of the observation (Julian year) used for proper motion correction.

        Returns
        -------
        sources : pandas.DataFrame
            Clean catalog
        """
        file_name = "%s/data/catalogs/ffi/%s/channel_%i_gaia_xmatch.csv" % (
            DATAOUTDIR,
            str(self.quarter),
            self.channel,
        )
        if os.path.isfile(file_name):
            print("Loading query from file...")
            print(file_name)
            sources = pd.read_csv(file_name).drop("Unnamed: 0", axis=1)
        else:
            # number of cells in the grid to divide the image
            nx = 4
            ny = 4
            stepx = int(self.ra_2d.shape[1] / nx)
            stepy = int(self.ra_2d.shape[0] / ny)
            sources = []
            for x in range(1, nx + 1):
                for y in range(1, ny + 1):
                    ra_cell = self.ra_2d[
                        (y - 1) * stepy : y * stepy, (x - 1) * stepx : x * stepx
                    ]
                    dec_cell = self.dec_2d[
                        (y - 1) * stepy : y * stepy, (x - 1) * stepx : x * stepx
                    ]

                    ra_q = np.mean(ra_cell)
                    dec_q = np.mean(dec_cell)
                    rad_q = np.hypot(
                        ra_cell - ra_cell.mean(), dec_cell - dec_cell.mean()
                    ).max()
                    print(
                        "Will do small queries query with this "
                        + "(ra, dec, radius, epoch): ",
                        ra_q,
                        dec_q,
                        rad_q,
                        epoch,
                    )
                    result = get_gaia_sources(
                        tuple([ra_q]),
                        tuple([dec_q]),
                        tuple([rad_q]),
                        magnitude_limit=18,
                        epoch=epoch,
                        dr=3,
                    )
                    sources.append(result)
            sources = pd.concat(sources, axis=0).drop_duplicates(subset=["designation"])
            print("Saving query to file...")
            print(file_name)
            columns = [
                "designation",
                "ra",
                "ra_error",
                "dec",
                "dec_error",
                "pmra",
                "pmdec",
                "parallax",
                "parallax_error",
                "phot_g_n_obs",
                "phot_g_mean_flux",
                "phot_g_mean_flux_error",
                "phot_g_mean_mag",
                "phot_bp_n_obs",
                "phot_bp_mean_flux",
                "phot_bp_mean_flux_error",
                "phot_bp_mean_mag",
                "phot_rp_n_obs",
                "phot_rp_mean_flux",
                "phot_rp_mean_flux_error",
                "phot_rp_mean_mag",
            ]
            sources = sources.loc[:, columns]

            if not os.path.isdir(
                "%s/data/catalogs/ffi/%s" % (DATAOUTDIR, str(self.quarter))
            ):
                os.makedirs("%s/data/catalogs/ffi/%s" % (DATAOUTDIR, str(self.quarter)))
            sources.to_csv(file_name)
        return sources

    def _clean_source_list(self, sources):
        """
        Function to clean surces from the catalog removing sources near the borders,
        with 10 pixel tolerance, and to remove blended sources (within 8")

        Parameters
        ----------
        sources : pandas.DataFrame
            Catalog with sources to be removed

        Returns
        -------
        sources : pandas.DataFrame
            Clean catalog
        """

        print("Cleaning sources table...")

        # find sources inside the image with 10 pix of inward tolerance
        inside = (
            (sources.row > 10)
            & (sources.row < 1014)
            & (sources.col > 10)
            & (sources.col < 1090)
        )
        sources = sources[inside].reset_index(drop=True)

        # find well separated sources
        s_coords = SkyCoord(sources.ra, sources.dec, unit=("deg"))
        midx, mdist = match_coordinates_3d(s_coords, s_coords, nthneighbor=2)[:2]
        # remove sources closer than 8" = 2 pix
        closest = mdist.arcsec < 8.0
        blocs = np.vstack([midx[closest], np.where(closest)[0]])
        bmags = np.vstack(
            [
                sources.phot_g_mean_mag[midx[closest]],
                sources.phot_g_mean_mag[np.where(closest)[0]],
            ]
        )
        faintest = [blocs[idx][s] for s, idx in enumerate(np.argmax(bmags, axis=0))]
        unresolved = np.in1d(np.arange(len(sources)), faintest)
        del s_coords, midx, mdist, closest, blocs, bmags

        sources = sources[~unresolved].reset_index(drop=True)

        return sources

    def _model_bkg(self, data, mask=None):
        """
        BkgZoomInterpolator:
        This class generates full-sized background and background RMS images
        from lower-resolution mesh images using the `~scipy.ndimage.zoom`
        (spline) interpolator.

        Parameters
        ----------
        data : numpy.ndarray
            Data arra with the pixel flux values.
        mask : numpy.ndarray
            Boolean array to mask pixels with sources.

        Returns
        -------
        background : numpy.ndarray
            Data array with background model
        """
        model = Background2D(
            data,
            mask=mask,
            box_size=(64, 50),
            filter_size=15,
            exclude_percentile=20,
            sigma_clip=SigmaClip(sigma=3.0, maxiters=5),
            bkg_estimator=MedianBackground(),
            interpolator=BkgZoomInterpolator(order=3),
        )

        return model.background

    def _saturated_pixels_mask(self, flux, column, row, saturation_limit=1.5e5):
        """
        Finds and removes saturated pixels, including bleed columns.

        Parameters
        ----------
        flux : numpu.ndarray
            Data array with pixel flux value
        column : numpy.ndarray
            Data array with pixel column value
        row : numpy.ndarray
            Data array with pixel row value
        saturation_limit : foat
            Saturation limit at which pixels are removed.

        Returns
        -------
        mask : numpy.ndarray
            Boolean mask with rejected pixels
        """
        # Which pixels are saturated
        # saturated = np.nanpercentile(flux, 99, axis=0)
        saturated = np.where((flux > saturation_limit).astype(float))[0]

        # Find bad pixels, including allowence for a bleed column.
        bad_pixels = np.vstack(
            [
                np.hstack([column[saturated] + idx for idx in np.arange(-3, 3)]),
                np.hstack([row[saturated] for idx in np.arange(-3, 3)]),
            ]
        ).T
        # Find unique row/column combinations
        bad_pixels = bad_pixels[
            np.unique(["".join(s) for s in bad_pixels.astype(str)], return_index=True)[
                1
            ]
        ]
        # Build a mask of saturated pixels
        m = np.zeros(len(column), bool)
        for p in bad_pixels:
            m |= (column == p[0]) & (row == p[1])
        return m

    def _mask_bright_sources(self, flux, column, row, sources, mag_limit=10):
        """
        Finds and removes halos produced by bright stars (<10 mag)

        Parameters
        ----------
        flux : numpu.ndarray
            Data array with pixel flux value
        column : numpy.ndarray
            Data array with pixel column value
        row : numpy.ndarray
            Data array with pixel row value
        sources : pandas.DataFrame
            Catalog wih observed sources in the image
        mag_limit : foat
            Magnitude limit at which bright sources are identified.

        Returns
        -------
        mask : numpy.ndarray
            Boolean mask with rejected pixels
        """
        bright_mask = sources["phot_g_mean_mag"] <= mag_limit
        mask_radius = 30  # Pixels

        mask = [
            np.hypot(column - s.col, row - s.row) < mask_radius
            for _, s in sources[bright_mask].iterrows()
        ]
        mask = np.array(mask).sum(axis=0) > 0

        return mask

    def _create_sparse(self):
        """
        Function to create sparse matrces (scipy.sparse.csr_matrix) for variables:
        dx, dy, dflux, dfluxerr, r, and phy
        This is extremelly necessary for FFI due to the large number of sources (~10k)
        and pixels (~1.1M). The sparse matrices contain the pixel data around the
        sources up to 7 pixels distance from the object location.
        """
        dx, dy, sparse_mask = [], [], []
        for i in tqdm(range(len(self.sources)), desc="Gaia sources"):
            dx_aux = self.col - self.sources["col"].iloc[i]
            dy_aux = self.row - self.sources["row"].iloc[i]
            near_mask = sparse.csr_matrix((np.abs(dx_aux) <= 7) & (np.abs(dy_aux) <= 7))

            dx.append(near_mask.multiply(dx_aux))
            dy.append(near_mask.multiply(dy_aux))
            sparse_mask.append(near_mask)

        del dx_aux, dy_aux, near_mask
        dx = sparse.vstack(dx, "csr")
        dy = sparse.vstack(dy, "csr")
        sparse_mask = sparse.vstack(sparse_mask, "csr")
        sparse_mask.eliminate_zeros()

        self.gf = self.sources["phot_g_mean_flux"].values
        self.dflux = sparse_mask.multiply(self.flux).tocsr()
        self.dflux_err = np.sqrt(np.abs(self.dflux))

        # eliminate leaked zero flux values in the sparse_mask
        self.sparse_mask = self.dflux.astype(bool)
        self.dx = self.sparse_mask.multiply(dx).tocsr()
        self.dy = self.sparse_mask.multiply(dy).tocsr()
        del dx, dy, sparse_mask

        # convertion to polar coordinates
        print("to polar coordinates...")
        nnz_inds = self.sparse_mask.nonzero()
        r_vals = np.hypot(self.dx.data, self.dy.data)
        phi_vals = np.arctan2(self.dy.data, self.dx.data)
        self.r = sparse.csr_matrix(
            (r_vals, (nnz_inds[0], nnz_inds[1])),
            shape=self.sparse_mask.shape,
            dtype=float,
        )
        self.phi = sparse.csr_matrix(
            (phi_vals, (nnz_inds[0], nnz_inds[1])),
            shape=self.sparse_mask.shape,
            dtype=float,
        )
        del r_vals, phi_vals, nnz_inds

        return

    def _get_source_mask(
        self,
        upper_radius_limit=7,
        lower_radius_limit=1.1,
        flux_cut_off=300,
        dm_type="rf-quadratic",
        plot=False,
    ):
        """
        Find the pixel mask that identifies pixels with contributions from ANY NUMBER
        of Sources.
        Fits a simple polynomial model to the log of the pixel flux values, in radial
        dimension and source flux, to find the optimum circular apertures for every
        source.

        Parameters
        ----------
        upper_radius_limit : float
            The radius limit at which we assume there is no flux from a source of any
            brightness (arcsec).
        lower_radius_limit : float
            The radius limit at which we assume there is flux from a source of any
            brightness (arcsec).
        flux_cut_off : float
            The flux at which we assume a source is too faint to model
        dm_type : string
            Type of design matrix to be used for modeling. Default is `rf-quadratic`,
            which is quadratic in both radius and flux.
        plot : bool
            Whether to show diagnostic plot. Default is False.
        """
        r = self.r
        mean_flux = self.dflux
        gf = self.gf

        nonz_idx = r.nonzero()
        rad_mask = r.data < upper_radius_limit
        temp_mask = sparse.csr_matrix(
            (r.data[rad_mask], (nonz_idx[0][rad_mask], nonz_idx[1][rad_mask])),
            shape=r.shape,
        ).astype(bool)
        temp_mask = temp_mask.multiply(temp_mask.sum(axis=0) == 1).tocsr()
        temp_mask.eliminate_zeros()

        with np.errstate(divide="ignore", invalid="ignore"):
            f = np.log10(temp_mask.astype(float).multiply(mean_flux).data)
        k = np.isfinite(f)
        f_mask = f[k]
        r_mask = temp_mask.astype(float).multiply(r).data[k]
        gf_mask = temp_mask.astype(float).multiply(gf[:, None]).data[k]
        k = np.isfinite(f_mask)

        A = make_A_edges(r_mask, np.log10(gf_mask), type=dm_type)

        for count in [0, 1, 2]:
            sigma_w_inv = A[k].T.dot(A[k])
            B = A[k].T.dot(f_mask[k])
            w = np.linalg.solve(sigma_w_inv, B)
            res = np.ma.masked_array(f_mask, ~k) - A.dot(w)
            k &= ~sigma_clip(res, sigma=3).mask

        test_f = np.linspace(
            np.log10(gf_mask.min()),
            np.log10(gf_mask.max()),
            100,
        )
        test_r = np.arange(lower_radius_limit, upper_radius_limit, 0.125)
        test_r2, test_f2 = np.meshgrid(test_r, test_f)

        test_A = make_A_edges(test_r2.ravel(), test_f2.ravel(), type=dm_type)
        test_val = test_A.dot(w).reshape(test_r2.shape)

        # find radius where flux > cut
        lr = np.zeros(len(test_f)) * np.nan
        for idx in range(len(test_f)):
            loc = np.where(10 ** test_val[idx] < flux_cut_off)[0]
            if len(loc) > 0:
                lr[idx] = test_r[loc[0]]

        ok = np.isfinite(lr)
        polifit_results = np.polyfit(test_f[ok], lr[ok], 2)
        source_radius_limit = np.polyval(polifit_results, np.log10(gf))
        source_radius_limit[
            source_radius_limit > upper_radius_limit
        ] = upper_radius_limit
        source_radius_limit[
            source_radius_limit < lower_radius_limit
        ] = lower_radius_limit

        self.radius = source_radius_limit + 0.5
        # self.source_mask = sparse.csr_matrix(self.r.value < self.radius[:, None])

        # remove pixels outside the radius limit
        source_mask = []
        for s in range(self.r.shape[0]):
            nonz_idx = self.r[s].nonzero()
            rad_mask = self.r[s].data < self.radius[s]
            aux = sparse.csr_matrix(
                (
                    self.r[s].data[rad_mask],
                    (nonz_idx[0][rad_mask], nonz_idx[1][rad_mask]),
                ),
                shape=self.r[s].shape,
            ).astype(bool)
            source_mask.append(aux)
        source_mask = sparse.vstack(source_mask, "csr")
        self.source_mask = source_mask

        if plot:
            fig, ax = plt.subplots(1, 2, figsize=(14, 5), facecolor="white")

            ax[0].scatter(r_mask, f_mask, s=0.4, c="k", alpha=0.5, label="Data")
            ax[0].scatter(
                r_mask[k],
                f_mask[k],
                s=0.4,
                c="g",
                alpha=0.5,
                label="Data clipped",
                rasterized=True,
            )
            ax[0].scatter(
                r_mask[k], A[k].dot(w), c="r", s=0.4, alpha=0.7, label="Model"
            )
            ax[0].set(
                xlabel=("Radius from Source [pix]"), ylabel=("log$_{10}$ Kepler Flux")
            )
            ax[0].legend(frameon=True, loc="upper right")

            im = ax[1].pcolormesh(
                test_f2,
                test_r2,
                10 ** test_val,
                vmin=0,
                vmax=500,
                cmap="viridis",
                shading="auto",
                rasterized=True,
            )
            line = np.polyval(np.polyfit(test_f[ok], lr[ok], 2), test_f)
            line[line > upper_radius_limit] = upper_radius_limit
            line[line < lower_radius_limit] = lower_radius_limit
            ax[1].plot(test_f, line, color="r", label="Best Fit PSF Edge")
            ax[1].legend(frameon=True, loc="upper left")
            cbar = plt.colorbar(im, ax=ax)
            cbar.set_label(r"PSF Flux [$e^-s^{-1}$]")

            ax[1].set(
                ylabel=("Radius from Source [pix]"),
                xlabel=("log$_{10}$ Source Flux"),
            )
            if not self.show:
                fig_name = "%s/data/figures/%s/channel_%02i_psf_edge_model_%s.png" % (
                    DATAOUTDIR,
                    str(self.quarter),
                    self.channel,
                    dm_type,
                )
                if not os.path.isdir("%s/data/figures/%i" % (DATAOUTDIR, self.quarter)):
                    os.makedirs("%s/data/figures/%i" % (DATAOUTDIR, self.quarter))

                plt.savefig(fig_name, format="png", bbox_inches="tight")
                plt.close()
                return

            plt.show()

        return

    def _get_uncontaminated_source_mask(self):
        """
        creates a mask of shape nsources x npixels where targets are not contaminated.
        This mask is used to select pixels to build the PSF model.
        """

        warnings.filterwarnings("ignore", category=sparse.SparseEfficiencyWarning)
        warnings.filterwarnings("ignore", category=RuntimeWarning)

        self.uncontaminated_source_mask = self.source_mask.multiply(
            self.source_mask.sum(axis=0) == 1
        ).tocsr()
        self.uncontaminated_source_mask.eliminate_zeros()

    # @profile
    def _build_prf_shape(self, n_r_knots=10, n_phi_knots=12, cut_r=1.5, flux_cut_off=1):
        """
        Builds a sparse model matrix of shape nsources x npixels to be used when
        fitting each source pixels to estimate its PSF photometry

        Parameters
        ----------
        n_r_knots : int
            Number of radial knots in the spline model.
        n_phi_knots : int
            Number of azimuthal knots in the spline model.
        cut_r : int
            Distance at which the spline Design matrix has only dependency in the
            radial axis.
        flux_cut_off: float
            The flux in COUNTS at which to stop evaluating the model.
        """

        flux_estimates = self.gf
        self.n_r_knots = n_r_knots
        self.n_phi_knots = n_phi_knots
        self.cut_r = cut_r

        # mean flux values using uncontaminated mask and normalized by flux estimations
        mean_f = np.log10(
            self.uncontaminated_source_mask.astype(float)
            .multiply(self.dflux)
            .multiply(1 / flux_estimates[:, None])
            .data
        )
        mean_f_err = np.abs(
            self.uncontaminated_source_mask.astype(float)
            .multiply(self.dflux_err)
            .multiply(1 / flux_estimates[:, None])
            .data
        )
        phi_b = self.uncontaminated_source_mask.multiply(self.phi).data
        r_b = self.uncontaminated_source_mask.multiply(self.r).data

        # build a design matrix A with b-splines basis in radius and angle axis.
        try:
            A = _make_A_polar(
                phi_b.ravel(),
                r_b.ravel(),
                cut_r=self.cut_r,
                rmin=self.rmin,
                rmax=self.rmax,
                n_r_knots=self.n_r_knots,
                n_phi_knots=self.n_phi_knots,
            )
        except ValueError:
            A = _make_A_polar(
                phi_b.ravel(),
                r_b.ravel(),
                cut_r=self.cut_r,
                rmin=self.rmin,
                rmax=np.percentile(r_b.ravel(), 98),
                n_r_knots=self.n_r_knots,
                n_phi_knots=self.n_phi_knots,
            )
            self.rmax = np.percentile(r_b.ravel(), 98)
        prior_sigma = np.ones(A.shape[1]) * 100
        prior_mu = np.zeros(A.shape[1]) - 10
        nan_mask = np.isfinite(mean_f.ravel())

        # we solve for A * psf_w = mean_f
        for count in [0, 1, 2]:
            psf_w = solve_linear_model(
                A,
                mean_f.ravel(),
                k=nan_mask,
                prior_mu=prior_mu,
                prior_sigma=prior_sigma,
                errors=False,
            )
            res = np.ma.masked_array(mean_f.ravel(), ~nan_mask) - A.dot(psf_w)
            nan_mask &= ~sigma_clip(res, sigma=3).mask

        self.psf_w = psf_w

        # We evaluate our DM and build PSF models per source
        self._get_mean_model()
        # mean_model = mean_model.multiply(1 / mean_model.sum(axis=1))

        #  re-estimate source flux (from CH updates)
        prior_mu = flux_estimates
        prior_sigma = np.ones(self.mean_model.shape[0]) * 10 * flux_estimates

        X = self.mean_model.copy().T

        fmean = self.uncontaminated_source_mask.astype(float).multiply(self.dflux).data
        femean = (
            self.uncontaminated_source_mask.astype(float).multiply(self.dflux_err).data
        )

        ws, werrs = solve_linear_model(
            X,
            fmean,
            y_err=femean,
            k=None,
            prior_mu=prior_mu,
            prior_sigma=prior_sigma,
            errors=True,
        )

        # Rebuild source mask
        ok = np.abs(ws - flux_estimates) / werrs > 3
        ok &= ((ws / flux_estimates) < 10) & ((flux_estimates / ws) < 10)
        ok &= ws > 10
        ok &= werrs > 0

        flux_estimates[ok] = ws[ok]

        self.source_mask = (
            self.mean_model.multiply(self.mean_model.T.dot(flux_estimates)).tocsr()
            > flux_cut_off
        )
        # rebuild uncontaminated_source_mask
        self._get_uncontaminated_source_mask()

        # set new rmax for spline basis
        self.rmax = np.minimum(
            self.rmax,
            np.percentile(self.uncontaminated_source_mask.multiply(self.r).data, 99.9),
        )
        self._get_mean_model()

        self.mean_flux = np.log10(
            self.uncontaminated_source_mask.astype(float)
            .multiply(self.dflux)
            .multiply(1 / flux_estimates[:, None])
            .data
        )

        print(
            "Total number of pixels data used for model fitting: ", self.mean_flux.shape
        )

        if self.save:
            self.save_model()

        if self.plot:
            self.plot_prf_shape()

        return

    def _get_mean_model(self):
        """
        Convenience function to make the scene PRF model
        """
        Ap = _make_A_polar(
            self.uncontaminated_source_mask.multiply(self.phi).data,
            self.uncontaminated_source_mask.multiply(self.r).data,
            rmin=self.rmin,
            rmax=self.rmax,
            cut_r=self.cut_r,
            n_r_knots=self.n_r_knots,
            n_phi_knots=self.n_phi_knots,
        )

        # And create a `mean_model` that has the psf model for all pixels with fluxes
        mean_model = sparse.csr_matrix(self.r.shape)
        m = 10 ** Ap.dot(self.psf_w)
        m[~np.isfinite(m)] = 0
        mean_model[self.uncontaminated_source_mask] = m
        mean_model.eliminate_zeros()
        self.mean_model = mean_model
        self.design_matrix = Ap

        return

    def build_prf_model(self, n_r_knots=5, n_phi_knots=15):
        """
        Function that creates a PRF shape using the sources. In combines all other
        helping functions that build the source mask, remove contaminated pixels,
        estimate PRF edges, and create the final PRF model.

        For details see:
            `self._create_sparse()`
            `self._get_source_mask()`
            `self._get_uncontaminated_source_mask()`
            `self._build_prf_shape()`

        Default parameters where used for Martinez-Palomera et al. 2021.

        Parameters
        ----------
        n_r_knots : int
            Number of radial knots in the spline model.
        n_phi_knots : int
            Number of azimuthal knots in the spline model.
        """
        psf._create_sparse()
        psf._get_source_mask(
            upper_radius_limit=5,
            lower_radius_limit=1.1,
            flux_cut_off=50,
            dm_type="rf-quadratic",
        )
        psf._get_uncontaminated_source_mask()
        psf._build_prf_shape(
            n_r_knots=n_r_knots, n_phi_knots=n_phi_knots, flux_cut_off=1
        )

    def save_model(self, path=None):
        """
        Function to save the PRF model weights, number of knots for r and phy, and
        rmin and rmax to re-build the Design Matrix.
        The file is a csv table, that contain a multi-index column table. Rows are each
        channel, and columns are:
            ["n_r_knots", "n_phi_knots", "rmin", "rmax", ...prf_ws...]
        This file can be loaded as:
            pd.read_csv(fname, index_col=0, header=[0, 1])

        Note: models with different number of knots lead to different number of weights,
        and ins necessary to create separete files to preserve the esctructure.

        Parameters
        ----------
        path : string
            Path of the file
        """
        if path is None:
            fname = "%s/data/ffi_prf_models_v0.1.1.csv" % (PACKAGEDIR)
        else:
            fname = path

        arr_to_save = np.array(
            [self.n_r_knots, self.n_phi_knots, self.rmin, self.rmax]
            + self.psf_w.tolist()
        )

        if not os.path.isfile(fname):
            df_dict = {
                self.quarter: pd.DataFrame(
                    np.atleast_2d(arr_to_save),
                    index=[self.channel],
                    columns=["n_r_knots", "n_phi_knots", "rmin", "rmax"]
                    + ["w%02i" % i for i in range(1, 1 + len(self.psf_w))],
                )
            }
            df = pd.concat(df_dict, axis=1, keys=df_dict.keys())
            df.to_csv(fname)

        else:
            df = pd.read_csv(fname, index_col=0, header=[0, 1])

            if str(self.quarter) in df.columns.levels[0]:
                if self.channel in df.index:
                    if (
                        int(df.loc[self.channel, (str(self.quarter), "n_r_knots")])
                        != self.n_r_knots
                        or int(df.loc[self.channel, (str(self.quarter), "n_phi_knots")])
                        != self.n_phi_knots
                    ):
                        raise ValueError(
                            "Number of knots for r or phi in the file does not"
                            + "matches the number used in the current model. "
                            + "Create a new file for current model."
                        )
                df.loc[self.channel, str(self.quarter)] = arr_to_save
            else:
                df_dict = {
                    self.quarter: pd.DataFrame(
                        np.atleast_2d(arr_to_save),
                        index=[self.channel],
                        columns=["n_r_knots", "n_phi_knots", "rmin", "rmax"]
                        + ["b%02i" % i for i in range(1, 1 + len(self.psf_w))],
                    )
                }
                df_new = pd.concat(df_dict, axis=1, keys=df_dict.keys())
                df = pd.concat([df, df_new], axis=1)

            # df.to_csv(fname)

        return

    def save_model_retro(self, path=None):
        """
        Function to save the PRF model as a pickle file.
        Depricated.

        Parameters
        ----------
        path : string
            Path of the file
        """
        model_data = dict(
            psf_w=self.psf_w,
            A=self.design_matrix,
            x_data=self.uncontaminated_source_mask.multiply(self.dx).data,
            y_data=self.uncontaminated_source_mask.multiply(self.dy).data,
            f_data=self.mean_flux,
            f_model=np.log10(self.mean_model.data),
            rmin=self.rmin,
            rmax=self.rmax,
            n_r_knots=self.n_r_knots,
            n_phi_knots=self.n_phi_knots,
        )

        if self.save:
            if path is None:
                output = "%s/data/models/%i/channel_%02i_psf_model.pkl" % (
                    DATAOUTDIR,
                    self.quarter,
                    self.channel,
                )
                if not os.path.isdir("%s/data/models/%i" % (DATAOUTDIR, self.quarter)):
                    os.makedirs("%s/data/models/%i" % (DATAOUTDIR, self.quarter))
            else:
                output = path
            with open(output, "wb") as file:
                pickle.dump(model_data, file)
        return

    def plot_prf_shape(self):
        """
        Function to plot the PRF model in Cartesian and Polar coordinates
        """
        ylim = self.uncontaminated_source_mask.multiply(self.r).data.max() * 1.1
        vmin = -3
        vmax = -0.5

        phy = self.uncontaminated_source_mask.multiply(self.phi).data
        r = self.uncontaminated_source_mask.multiply(self.r).data
        x = self.uncontaminated_source_mask.multiply(self.dx).data
        y = self.uncontaminated_source_mask.multiply(self.dy).data

        fig, ax = plt.subplots(2, 2, figsize=(12, 8))
        ax[0, 0].set_title("Mean flux")
        cax = ax[0, 0].scatter(
            phy,
            r,
            c=self.mean_flux,
            marker=".",
            s=2,
            vmin=vmin,
            vmax=vmax,
            rasterized=True,
        )
        ax[0, 0].set_ylim(0, ylim)
        fig.colorbar(cax, ax=ax[0, 0])
        ax[0, 0].set_ylabel(r"$r$ [pixels]")
        ax[0, 0].set_xlabel(r"$\phi$ [rad]")

        ax[0, 1].set_title("Average PSF Model")
        cax = cax = ax[0, 1].scatter(
            phy,
            r,
            c=np.log10(self.mean_model.data),
            marker=".",
            s=2,
            vmin=vmin,
            vmax=vmax,
            rasterized=True,
        )
        ax[0, 1].set_ylim(0, ylim)
        fig.colorbar(cax, ax=ax[0, 1])
        ax[0, 1].set_xlabel(r"$\phi$ [rad]")

        cax = ax[1, 0].scatter(
            x,
            y,
            c=self.mean_flux,
            marker=".",
            s=2,
            vmin=vmin,
            vmax=vmax,
            rasterized=True,
        )
        fig.colorbar(cax, ax=ax[1, 0])
        ax[1, 0].set_ylabel("dy")
        ax[1, 0].set_xlabel("dx")

        cax = ax[1, 1].scatter(
            x,
            y,
            c=np.log10(self.mean_model.data),
            marker=".",
            s=2,
            vmin=vmin,
            vmax=vmax,
            rasterized=True,
        )
        fig.colorbar(cax, ax=ax[1, 1])
        ax[1, 1].set_xlabel("dx")

        if not self.show:
            fig_name = "%s/data/figures/%s/channel_%02i_psf_model.png" % (
                DATAOUTDIR,
                str(self.quarter),
                self.channel,
            )
            if not os.path.isdir(
                "%s/data/figures/%s" % (DATAOUTDIR, str(self.quarter))
            ):
                os.makedirs("%s/data/figures/%s" % (DATAOUTDIR, str(self.quarter)))
            plt.savefig(fig_name, format="png", bbox_inches="tight")
            plt.close()
        else:
            plt.show()

    # @profile
    def fit_model(self):
        """
        Function to fit the PRF model and do LFD photometry for the sources observed in
        the FFIs.
        """
        prior_mu = self.gf
        prior_sigma = np.ones(self.mean_model.shape[0]) * 5 * np.abs(self.gf) ** 0.5

        X = self.mean_model.copy()
        X = X.T
        f = self.flux
        fe = self.flux_err

        self.ws, self.werrs = solve_linear_model(
            X, f, y_err=fe, prior_mu=prior_mu, prior_sigma=prior_sigma, errors=True
        )
        self.model_flux = X.dot(self.ws)

        nodata = np.asarray(self.source_mask.sum(axis=1))[:, 0] == 0
        # These sources are poorly estimated
        nodata |= (self.mean_model.max(axis=1) > 1).toarray()[:, 0]
        self.ws[nodata] *= np.nan
        self.werrs[nodata] *= np.nan

        return

    def save_catalog(self):
        """
        Function to save the Photometry Catalog of FFI sources
        """
        df = pd.DataFrame(
            [
                self.sources.designation,
                self.sources.ra,
                self.sources.dec,
                self.sources.col,
                self.sources.row,
                self.ws,
                self.werrs,
            ],
            index=["Gaia_source_id", "RA", "DEC", "Column", "Row", "Flux", "Flux_err"],
        ).T

        if not os.path.isdir("%s/data/catalogs/ffi/source_catalog/" % (DATAOUTDIR)):
            os.makedirs("%s/data/catalogs/ffi/source_catalog/" % (DATAOUTDIR))
        df.to_csv(
            "%s/data/catalogs/ffi/source_catalog/channel_%s_source_catalog_mjd_%s.csv"
            % (DATAOUTDIR, self.channel, str(self.hdr["MJDSTART"]))
        )

    def plot_image(self, ax=None, sources=False):
        """
        Function to plot the Full Frame Image and the Gaia Sources

        Parameters
        ----------
        ax : matplotlib.axes
            Matlotlib axis can be provided, if not one will be created and returned
        sources : boolean
            Whether to overplot or not the source catalog

        Returns
        -------
        ax : matplotlib.axes
            Matlotlib axis with the figure
        """
        if ax is None:
            fig, ax = plt.subplots(1, figsize=(10, 10))
        ax = plt.subplot(projection=self.wcs)
        im = ax.imshow(
            self.flux_2d,
            cmap=plt.cm.viridis,
            origin="lower",
            norm=colors.SymLogNorm(linthresh=200, vmin=0, vmax=2000, base=10),
            rasterized=True,
        )
        plt.colorbar(im, ax=ax, label=r"Flux ($e^{-}s^{-1}$)", fraction=0.042)

        ax.set_title("FFI Ch %i" % (self.channel))
        ax.set_xlabel("R.A. [hh:mm]")
        ax.set_ylabel("Decl. [deg]")
        ax.grid(color="white", ls="solid")
        ax.set_aspect("equal", adjustable="box")

        if sources:
            ax.scatter(
                self.sources.col,
                self.sources.row,
                facecolors="none",
                edgecolors="r",
                linewidths=0.5,
                alpha=0.9,
            )

        if self.save:
            fig_name = "%s/data/figures/%s/channel_%02i_ffi_image.png" % (
                DATAOUTDIR,
                str(self.quarter),
                self.channel,
            )
            if not os.path.isdir(
                "%s/data/figures/%s" % (DATAOUTDIR, str(self.quarter))
            ):
                os.makedirs("%s/data/figures/%s" % (DATAOUTDIR, str(self.quarter)))
            plt.savefig(fig_name, format="png", bbox_inches="tight")

        return ax

    def plot_pixel_masks(self, ax=None):
        """
        Function to plot the mask used to reject saturated and bright pixels

        Parameters
        ----------
        ax : matplotlib.axes
            Matlotlib axis can be provided, if not one will be created and returned

        Returns
        -------
        ax : matplotlib.axes
            Matlotlib axis with the figure
        """
        if ax is None:
            fig, ax = plt.subplots(1, figsize=(10, 10))
        ax.scatter(
            self.col_2d.ravel()[self.non_sat_mask][~self.bright_mask],
            self.row_2d.ravel()[self.non_sat_mask][~self.bright_mask],
            c="r",
            marker=".",
            label="bright",
        )
        ax.scatter(
            self.col_2d.ravel()[~self.non_sat_mask],
            self.row_2d.ravel()[~self.non_sat_mask],
            c="y",
            marker=".",
            label="saturated",
        )
        ax.legend(loc="best")

        ax.set_xlabel("Column Pixel Number")
        ax.set_ylabel("Row Pixel Number")
        ax.set_title("Pixel Mask")

        if self.save:
            fig_name = "%s/data/figures/%s/channel_%02i_ffi_pixel_mask.png" % (
                DATAOUTDIR,
                str(self.quarter),
                self.channel,
            )
            if not os.path.isdir(
                "%s/data/figures/%s" % (DATAOUTDIR, str(self.quarter))
            ):
                os.makedirs("%s/data/figures/%s" % (DATAOUTDIR, str(self.quarter)))
            plt.savefig(fig_name, format="png", bbox_inches="tight")

        return ax


'''
This algorithm is a standard TD3 algorithm.
Original paper: https://arxiv.org/abs/1802.09477.
'''
import pickle
import numpy as np

import torch
import torch.nn as nn

from TD3_based_DRL.priority_replay import Memory
from TD3_based_DRL.network_model import Actor,Critic
from TD3_based_DRL.util import hard_update, soft_update

seed = 2
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

MEMORY_CAPACITY = 38400
BATCH_SIZE = 128
GAMMA = 0.95
LR_C = 0.0005
LR_A = 0.0002
LR_I = 0.01
TAU = 0.001
POLICY_NOSIE = 0.2
POLICY_FREQ = 1
NOISE_CLIP = 0.5

class DRL:

    def __init__(self,action_dim,state_dim, LR_C = LR_C, LR_A = LR_A):
        self.use_cuda = True

        self.state_dim = state_dim[0] * state_dim[1]
        self.state_dim_width = state_dim[0]
        self.state_dim_height = state_dim[1]
        self.action_dim = action_dim
        self.batch_size = BATCH_SIZE
        self.gamma = GAMMA
        self.tau = TAU
        self.policy_noise = POLICY_NOSIE
        self.noise_clip = NOISE_CLIP
        self.policy_freq = POLICY_FREQ
        self.itera = 0

        self.pointer = 0
        self.memory = Memory(MEMORY_CAPACITY)

        self.actor = Actor(self.state_dim,self.action_dim).to(self.device)
        self.actor_target = Actor(self.state_dim,self.action_dim).to(self.device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),LR_A)

        self.critic = Critic(self.state_dim,self.action_dim).to(self.device)
        self.critic_target = Critic(self.state_dim,self.action_dim).to(self.device)
        self.critic_optimizers = torch.optim.Adam(self.critic.parameters(),LR_C)

        hard_update(self.actor_target,self.actor)
        hard_update(self.critic_target,self.critic)

    def learn(self, batch_size=BATCH_SIZE, epoch=0):

        ## batched state, batched action, batched reward, batched next state
        bs, ba, ba_e, bi, br, bs_, tree_idx, ISweight = self.retrive(batch_size)
        bs = torch.tensor(bs, dtype=torch.float).reshape(batch_size, self.state_dim_height, self.state_dim_width).to(self.device)
        ba = torch.tensor(ba, dtype=torch.float).to(self.device)
        br = torch.tensor(br, dtype=torch.float).to(self.device)
        bs_ = torch.tensor(bs_, dtype=torch.float).reshape(batch_size, self.state_dim_height, self.state_dim_width).to(self.device)

        # initialize the loss variables
        loss_c, loss_a = 0, 0

        ## calculate the predicted values of the critic
        with torch.no_grad():
            noise = (torch.randn_like(ba) * self.policy_noise).clamp(0, 1)
            a_ = (self.actor_target(bs_).detach() + noise).clamp(0, 1)
            target_q1, target_q2 = self.critic_target([bs_,a_])
            target_q1 = target_q1.detach()
            target_q2 = target_q2.detach()
            target_q = torch.min(target_q1,target_q2)
            y_expected = br + self.gamma * target_q
        y_predicted1, y_predicted2 = self.critic.forward([bs,ba])

        ## update the critic
        critic_loss = nn.MSELoss()
        loss_critic = critic_loss(y_predicted1,y_expected)+critic_loss(y_predicted2,y_expected)
        self.critic_optimizers.zero_grad()
        loss_critic.backward()
        self.critic_optimizers.step()

        ## update the actor
        if self.itera % self.policy_freq == 0:

            pred_a = self.actor.forward(bs)
            loss_actor = (-self.critic.forward([bs,pred_a])[0])

            self.actor_optimizer.zero_grad()
            loss_actor.mean().backward()
            self.actor_optimizer.step()

            soft_update(self.actor_target,self.actor,self.tau)
            soft_update(self.critic_target,self.critic,self.tau)

            loss_a = loss_actor.mean().item()

        loss_c = loss_critic.mean().item()

        self.itera += 1

        self.memory.batch_update(tree_idx, abs(errors.detach().cpu().numpy()) )

        return loss_c, loss_a

    def choose_action(self,state):

        state = torch.tensor(state,dtype=torch.float).reshape(self.state_dim_height, self.state_dim_width).to(self.device)
        state = state.unsqueeze(0)

        action = self.actor.forward(state).detach()
        action = action.squeeze(0).cpu().numpy()
        action = np.clip(action,-1, 1)

        return action

    '''
    To conveniently implement all the algorithms in the same training script, the memory buffer here
    keeps the same shape as other human-guidance DRL algorithms. However, in this vanilla TD3 algorithm,
    the human-guidance-related indicators, i.e., action from the human expert, and the intervention signal,
    are not available in fact.
    '''
    def store_transition(self, s, a, a_e, i, r, s_):

        transition = np.hstack((s, a, a_e, i, r, s_))
        self.memory.store(transition)
        self.pointer += 1

    def retrive(self, batch_size):

        tree_index, bt, ISWeight = self.memory.sample(batch_size)
        bs = bt[:, :self.state_dim]
        ba = bt[:, self.state_dim: self.state_dim + self.action_dim]
        ba_e = bt[:, self.state_dim + self.action_dim: self.state_dim + self.action_dim + self.action_dim]
        bi = bt[:, -self.state_dim - 2: -self.state_dim - 1]
        br = bt[:, -self.state_dim - 1: -self.state_dim]
        bs_ = bt[:, -self.state_dim:]

        return bs, ba, ba_e, bi, br, bs_, tree_index, ISWeight

    def memory_save(self):

        per = open("memory.pkl", 'wb')
        str = pickle.dumps(self.memory)
        per.write(str)
        per.close()

    def memory_load(self):

        with open("memory.pkl",'rb') as file:
            self.memory  = pickle.loads(file.read())

    def load_model(self, output):
        if output is None: return
        self.actor.load_state_dict(torch.load('{}/actor.pkl'.format(output)))
        self.critic.load_state_dict(torch.load('{}/critic.pkl'.format(output)))

    def save_model(self, output):
        torch.save(self.actor.state_dict(), '{}/actor.pkl'.format(output))
        torch.save(self.critic.state_dict(), '{}/critic.pkl'.format(output))

    def save(self, log_dir, epoch):
        state = {'actor':self.actor.state_dict(), 'actor_target':self.actor_target.state_dict(),
                 'actor_optimizer':self.actor_optimizer.state_dict(),
                 'critic':self.critic.state_dict(), 'critic_target':self.critic_target.state_dict(),
                 'critic_optimizers':self.critic_optimizers.state_dict(),
                 'epoch':epoch}
        torch.save(state, log_dir)

    def load(self, log_dir):
        checkpoint = torch.load(log_dir)
        self.actor.load_state_dict(checkpoint['actor'])
        self.actor_target.load_state_dict(checkpoint['actor_target'])
        self.actor_optimizer.load_state_dict(checkpoint['actor_optimizer'])
        self.critic.load_state_dict(checkpoint['critic'])
        self.critic_target.load_state_dict(checkpoint['critic_target'])
        self.critic_optimizers.load_state_dict(checkpoint['critic_optimizers'])