/* The MIT License

   Copyright (c) 2013-2015 Genome Research Ltd.

   Author: Petr Danecek <pd3@sanger.ac.uk>

   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   copies of the Software, and to permit persons to whom the Software is
   furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in
   all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
   THE SOFTWARE.

 */

#include <stdio.h>
#include <stdlib.h>
#include <getopt.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_multifit_nlin.h>
#include <htslib/vcf.h>
#include <htslib/synced_bcf_reader.h>
#include "bcftools.h"
#include "peakfit.h"

typedef struct
{
    int nvals;          // all values, including RR,AA peaks
    double *xvals;      // pointer to args_t.xvals
    double *yvals;
    int copy_number;    // heuristics to skip futile CN1 fits when no het peak is detected
    int irr, ira, iaa;  // chop off RR and AA peaks
    char *chr;
}
dist_t;

typedef struct
{
    int ndist, nbins, ra_rr_scaling, smooth;
    double *xvals;
    dist_t *dist;
    char **argv, *output_dir;
    double fit_th, peak_symmetry, cn_penalty, min_peak_size, min_fraction;
    int argc, plot, verbose, regions_is_file, targets_is_file, include_aa, force_cn;
    char *dat_fname, *fname, *regions_list, *targets_list, *sample;
    FILE *dat_fp;
}
args_t;

FILE *open_file(char **fname, const char *mode, const char *fmt, ...);

static void init_dist(args_t *args, dist_t *dist, int verbose)
{
    // isolate RR and AA peaks and rescale so that they are comparable to hets
    int i, irr, iaa, n = dist->nvals;

    // smooth the distribution, this is just to find the peaks
    double *tmp = (double*) malloc(sizeof(double)*n);
    int win  = args->smooth ? fabs(args->smooth)*2 + 1 : 7;   // must be an odd number
    int hwin = win/2;
    double avg = tmp[0] = dist->yvals[0];
    for (i=1; i<hwin; i++)
    {
        avg += dist->yvals[2*i-1];
        tmp[i] = avg/(2*i+1);
    }
    avg = 0;
    for (i=0; i<n; i++)
    {
        avg += dist->yvals[i];
        if ( i>=win-1 )
        {
            tmp[i-hwin] = avg/win;
            avg -= dist->yvals[i-win+1];
        }
    }
    for (i=n-hwin; i<n; i++)
    {
        avg -= dist->yvals[i-hwin];
        hwin--;
        tmp[i] = avg/(2*hwin+1);
        avg -= dist->yvals[i-hwin];
    }

    // find the extremes; first a simple approach: find a gap
    for (irr=0,i=0; i<n/2; i++) if ( tmp[i] < tmp[irr] ) irr = i;
    for (iaa=n-1,i=n-1; i>=n/2; i--) if ( tmp[i] < tmp[iaa] ) iaa = i;
    irr += win*0.5;
    iaa += win*0.5;
    if ( iaa>=n ) iaa = n-1;
    if ( irr>=iaa ) error("FIXME: oops, dist normalization failed for %s: %d vs %d\n", dist->chr,irr,iaa); // we may need to be smarter
    if ( args->smooth>0 ) for (i=0; i<n; i++) dist->yvals[i] = tmp[i];
    free(tmp);

    // clean the data: the AA peak is occasionally not centered at 1.0 but is closer to the center, chop off
    int imax_aa = iaa;
    for (i=iaa; i<n; i++)
        if ( dist->yvals[imax_aa] < dist->yvals[i] ) imax_aa = i;
    dist->nvals = imax_aa+1;
    if ( iaa>=dist->nvals ) iaa = dist->nvals-1;

    // find the maximum and scale the peaks (first draft: no attempt to join the segments smootly)
    double max_rr = 0, max_aa = 0, max_ra = 0, srr = 0, saa = 0, sra = 0;
    for (i=0; i<irr; i++)
    {
        srr += dist->yvals[i];
        if ( max_rr < dist->yvals[i] ) max_rr = dist->yvals[i];
    }
    for (i=irr; i<=iaa; i++)
    {
        sra += dist->yvals[i];
        if ( max_ra < dist->yvals[i] ) max_ra = dist->yvals[i];
    }
    for (i=iaa+1; i<n; i++)
    {
        saa += dist->yvals[i];
        if ( max_aa < dist->yvals[i] ) max_aa = dist->yvals[i];
    }

    // Does the het peak exist at all? Usually the numbers are as follows:
    //  1:  cn=0     ra/rr=0.205730      aa/ra=0.674922      nra=7066
    //  20: cn=0     ra/rr=0.258019      aa/ra=0.548929      nra=2381
    //  X:  cn=0     ra/rr=0.005976      aa/ra=44.116667     nra=60
    //  Y:  cn=0     ra/rr=0.008316      aa/ra=7.250000      nra=12
    //  MT: cn=0     ra/rr=0.013699      aa/ra=0.666667      nra=3

    if ( !args->ra_rr_scaling ) max_ra = max_aa = max_rr;
    if ( !sra || (sra/srr<0.1 && saa/sra>1.0) ) // too few hets, CN1
    {
        max_ra = max_aa;
        dist->copy_number = 1;
    }
    else if ( sra/srr<0.1 || saa/sra>1.0 )
    {
        max_ra = max_aa;
        dist->copy_number = -1;     // unknown copy number
    }
    if ( max_rr ) for (i=0; i<irr; i++) dist->yvals[i] /= max_rr;
    if ( max_ra ) for (i=irr; i<=iaa; i++) dist->yvals[i] /= max_ra;
    if ( max_aa ) for (i=iaa+1; i<n; i++) dist->yvals[i] /= max_aa;

    dist->irr = irr;
    dist->iaa = iaa;
    dist->ira = n*0.5;

    if ( verbose )
        fprintf(stderr,"%s:\t irr,ira,iaa=%.2f,%.2f,%.2f \t cn=%2d \t ra/rr=%f \t aa/ra=%f \t nra=%d\n",
            dist->chr, dist->xvals[irr],dist->xvals[dist->ira],dist->xvals[iaa],
            dist->copy_number,sra/srr,saa/sra, (int)sra);
}

static void init_data(args_t *args)
{
    bcf_srs_t *files = bcf_sr_init();
    if ( args->regions_list )
    {
        if ( bcf_sr_set_regions(files, args->regions_list, args->regions_is_file)<0 )
            error("Failed to read the regions: %s\n", args->regions_list);
    }
    if ( args->targets_list )
    {
        if ( bcf_sr_set_targets(files, args->targets_list, args->targets_is_file, 0)<0 )
            error("Failed to read the targets: %s\n", args->targets_list);
    }
    if ( !bcf_sr_add_reader(files, args->fname) ) error("Failed to open %s: %s\n", args->fname,bcf_sr_strerror(files->errnum));
    bcf_hdr_t *hdr = files->readers[0].header;
    if ( !args->sample )
    {
        if ( bcf_hdr_nsamples(hdr)>1 ) error("Missing the option -s, --sample\n");
        args->sample = hdr->samples[0];
    }
    else if ( bcf_hdr_id2int(hdr,BCF_DT_SAMPLE,args->sample)<0 ) error("No such sample: %s\n", args->sample);
    int ret = bcf_hdr_set_samples(hdr, args->sample, 0);
    if ( ret<0 ) error("Error setting the sample: %s\n", args->sample);

    if ( !bcf_hdr_idinfo_exists(hdr,BCF_HL_FMT,bcf_hdr_id2int(hdr,BCF_DT_ID,"BAF")) )
        error("The tag FORMAT/BAF is not present in the VCF: %s\n", args->fname);

    int i;
    args->xvals = (double*) calloc(args->nbins,sizeof(double));
    for (i=0; i<args->nbins; i++) args->xvals[i] = 1.0*i/(args->nbins-1);

    // collect BAF distributions for all chromosomes
    int idist = -1, nbaf = 0, nprocessed = 0, ntotal = 0, prev_chr = -1;
    float *baf = NULL;
    while ( bcf_sr_next_line(files) )
    {
        ntotal++;

        bcf1_t *line = bcf_sr_get_line(files,0);
        if ( bcf_get_format_float(hdr,line,"BAF",&baf,&nbaf) != 1 ) continue;
        if ( bcf_float_is_missing(baf[0]) ) continue;

        nprocessed++;

        if ( prev_chr==-1 || prev_chr!=line->rid )
        {
            // new chromosome
            idist = args->ndist++;
            args->dist = (dist_t*) realloc(args->dist, sizeof(dist_t)*args->ndist);
            memset(&args->dist[idist],0,sizeof(dist_t));
            args->dist[idist].chr   = strdup(bcf_seqname(hdr,line));
            args->dist[idist].yvals = (double*) calloc(args->nbins,sizeof(double));
            args->dist[idist].xvals = args->xvals;
            args->dist[idist].nvals = args->nbins;
            prev_chr = line->rid;
        }
        int bin = baf[0]*(args->nbins-1);
        args->dist[idist].yvals[bin]++;   // the distribution
    }
    free(baf);
    bcf_sr_destroy(files);

    for (idist=0; idist<args->ndist; idist++)
    {
        #if 0
            int j;
            for (j=0; j<args->nbins; j++)
            {
                double x = args->dist[idist].xvals[j];
                args->dist[idist].yvals[j] = exp(-(x-0.5)*(x-0.5)/1e-3);
            }
        #endif
        init_dist(args, &args->dist[idist],args->verbose);
    }

    args->dat_fp = open_file(&args->dat_fname,"w","%s/dist.dat", args->output_dir);
    fprintf(args->dat_fp, "# This file was produced by: bcftools polysomy(%s+htslib-%s), the command line was:\n", bcftools_version(),hts_version());
    fprintf(args->dat_fp, "# \t bcftools %s ", args->argv[0]);
    for (i=1; i<args->argc; i++)
        fprintf(args->dat_fp, " %s",args->argv[i]);
    fprintf(args->dat_fp,"\n#\n");
    fprintf(args->dat_fp,"# DIST\t[2]Chrom\t[3]BAF\t[4]Normalized Count\n");
    fprintf(args->dat_fp,"# FIT\t[2]Goodness of Fit\t[3]iFrom\t[4]iTo\t[5]The Fitted Function\n");
    fprintf(args->dat_fp,"# CN\t[2]Chrom\t[3]Estimated Copy Number\t[4]Absolute fit deviation\n");

    char *fname = NULL;
    FILE *fp = open_file(&fname,"w","%s/dist.py", args->output_dir);
//-------- matplotlib script --------------
    fprintf(fp,
        "#!/usr/bin/env python\n"
        "#\n"
        "import matplotlib as mpl\n"
        "mpl.use('Agg')\n"
        "import matplotlib.pyplot as plt\n"
        "import csv,sys,argparse\n"
        "from math import exp\n"
        "\n"
        "outdir = '%s'\n"
        "\n"
        "def read_dat(dat,fit,cn):\n"
        "   csv.register_dialect('tab', delimiter='\t', quoting=csv.QUOTE_NONE)\n"
        "   with open(outdir+'/dist.dat', 'rb') as f:\n"
        "      reader = csv.reader(f, 'tab')\n"
        "      for row in reader:\n"
        "          if row[0][0]=='#': continue\n"
        "          type = row[0]\n"
        "          chr  = row[1]\n"
        "          if type=='DIST':\n"
        "              if chr not in dat: dat[chr] = []\n"
        "              dat[chr].append(row)\n"
        "          elif type=='FIT':\n"
        "              if chr not in fit: fit[chr] = []\n"
        "              fit[chr].append(row)\n"
        "          elif type=='CN':\n"
        "              cn[chr] = row[2]\n"
        "\n"
        "def plot_dist(dat,fit,chr):\n"
        "   fig, ax = plt.subplots(1, 1, figsize=(7,5))\n"
        "   ax.plot([x[2] for x in dat[chr]],[x[3] for x in dat[chr]],'k-',label='Distribution')\n"
        "   if chr in fit:\n"
        "       for i in range(len(fit[chr])):\n"
        "           pfit = fit[chr][i]\n"
        "           exec('def xfit(x): return '+pfit[5])\n"
        "           istart = int(pfit[3])\n"
        "           iend   = int(pfit[4])+1\n"
        "           vals   = dat[chr][istart:iend]\n"
        "           args   = {}\n"
        "           if i==0: args = {'label':'Target to Fit'}\n"
        "           ax.plot([x[2] for x in vals],[x[3] for x in vals],'r-',**args)\n"
        "           if i==0: args = {'label':'Best Fit'}\n"
        "           ax.plot([x[2] for x in vals],[xfit(float(x[2])) for x in vals],'g-',**args)\n"
        "   ax.set_title('BAF distribution, chr'+chr)\n"
        "   ax.set_xlabel('BAF')\n"
        "   ax.set_ylabel('Frequency')\n"
        "   ax.legend(loc='best',prop={'size':7},frameon=False)\n"
        "   plt.savefig(outdir+'/dist.chr'+chr+'.png')\n"
        "   plt.close()\n"
        "\n"
        "def plot_copy_number(cn):\n"
        "   fig, ax = plt.subplots(1, 1, figsize=(7,5))\n"
        "   xlabels = sorted(cn.keys())\n"
        "   xvals = range(len(xlabels))\n"
        "   yvals = [float(cn[x]) for x in xlabels]\n"
        "   ax.plot(xvals,yvals,'o',color='red')\n"
        "   for i in range(len(xvals)):\n"
        "       if yvals[i]==-1: ax.annotate('?', xy=(xvals[i],0.5),va='center',ha='center',color='red',fontweight='bold')\n"
        "   ax.tick_params(axis='both', which='major', labelsize=9)\n"
        "   ax.set_xticks(xvals)\n"
        "   ax.set_xticklabels(xlabels,rotation=45)\n"
        "   ax.set_xlim(-1,len(xlabels))\n"
        "   ax.set_ylim(0,5.0)\n"
        "   ax.set_yticks([1.0,2.0,3.0,4.0])\n"
        "   ax.set_xlabel('Chromosome')\n"
        "   ax.set_ylabel('Copy Number')\n"
        "   plt.savefig(outdir+'/copy-number.png')\n"
        "   plt.close()\n"
        "\n"
        "class myParser(argparse.ArgumentParser):\n"
        "   def error(self, message):\n"
        "       self.print_help()\n"
        "       sys.stderr.write('error: %%s\\n' %% message)\n"
        "       sys.exit(2)\n"
        "\n"
        "def main():\n"
        "   parser = myParser()\n"
        "   parser.add_argument('-a', '--all', action='store_true', help='Create all plots')\n"
        "   parser.add_argument('-c', '--copy-number', action='store_true', help='Create copy-number plot')\n"
        "   parser.add_argument('-d', '--distrib', metavar='CHR', help='Plot BAF distribution of a single chromosome')\n"
        "   args = parser.parse_args()\n"
        "   dat = {}; fit = {}; cn = {}\n"
        "   read_dat(dat,fit,cn)\n"
        "   if args.distrib!=None:\n"
        "       plot_dist(dat,fit,args.distrib)\n"
        "   if args.all:\n"
        "       for chr in dat: plot_dist(dat,fit,chr)\n"
        "       plot_copy_number(cn)\n"
        "   elif args.copy_number:\n"
        "       plot_copy_number(cn)\n"
        "   else:\n"
        "       for chr in dat: plot_dist(dat,fit,chr)\n"
        "\n"
        "if __name__ == '__main__':\n"
        "   main()\n",
        args->output_dir);
//---------------------------------------
    chmod(fname, S_IWUSR|S_IRUSR|S_IRGRP|S_IROTH|S_IXUSR|S_IXGRP|S_IXOTH);
    free(fname);
    fclose(fp);
}

static void destroy_data(args_t *args)
{
    int i;
    for (i=0; i<args->ndist; i++)
    {
        free(args->dist[i].chr);
        free(args->dist[i].yvals);
    }
    free(args->dist);
    free(args->xvals);
    free(args->dat_fname);
    fclose(args->dat_fp);
}

static void save_dist(args_t *args, dist_t *dist)
{
    int i;
    for (i=0; i<args->nbins; i++)
        fprintf(args->dat_fp,"DIST\t%s\t%f\t%f\n",dist->chr,dist->xvals[i],dist->yvals[i]);
}
static void fit_curves(args_t *args)
{
    peakfit_t *pkf = peakfit_init();
    peakfit_verbose(pkf,args->verbose);

    int i, nmc = 50;
    for (i=0; i<args->ndist; i++)
    {
        dist_t *dist = &args->dist[i];
        save_dist(args, &args->dist[i]);

        if ( dist->copy_number!=0 )
        {
            fprintf(args->dat_fp,"CN\t%s\t%.2f\n", dist->chr,(float)dist->copy_number);
            continue;
        }

        if ( args->verbose )
            fprintf(stderr,"%s:\n", dist->chr);

        int nrr_aa  = dist->iaa - dist->irr + 1;
        int nrr_ra  = dist->ira - dist->irr + 1;
        int naa_max = dist->nvals - dist->iaa;
        double xrr  = dist->xvals[dist->irr], *xrr_vals = &dist->xvals[dist->irr], *yrr_vals = &dist->yvals[dist->irr];
        double xaa  = dist->xvals[dist->iaa], *xaa_vals = &dist->xvals[dist->iaa], *yaa_vals = &dist->yvals[dist->iaa];
        double xra  = dist->xvals[dist->ira];
        double xmax = dist->xvals[dist->nvals-1];

        // CN2
        double cn2aa_fit = 0, cn2ra_fit, cn2_fit;
        char *cn2aa_func = 0, *cn2ra_func;
        double cn2aa_params[3] = {1,1,1} ,cn2ra_params[3];
        if ( args->include_aa )
        {
            peakfit_reset(pkf);
            peakfit_add_exp(pkf, 1.0,1.0,0.2, 5);
            peakfit_set_mc(pkf, 0.01,0.3,2,nmc);
            peakfit_set_mc(pkf, 0.05,1.0,0,nmc);
            cn2aa_fit  = peakfit_run(pkf, naa_max, xaa_vals, yaa_vals);
            cn2aa_func = strdup(peakfit_sprint_func(pkf));
            peakfit_get_params(pkf,0,cn2aa_params,3);
        }
        peakfit_reset(pkf);
        peakfit_add_bounded_gaussian(pkf, 1.0,0.5,0.03, 0.45,0.55, 7);
        peakfit_set_mc(pkf, 0.01,0.3,2,nmc);
        peakfit_set_mc(pkf, 0.05,1.0,0,nmc);
        cn2ra_fit  = peakfit_run(pkf, nrr_aa,xrr_vals,yrr_vals);
        cn2ra_func = strdup(peakfit_sprint_func(pkf));
        cn2_fit    = cn2ra_fit + cn2aa_fit;
        peakfit_get_params(pkf,0,cn2ra_params,3);

        // CN3: fit two peaks, then enforce the symmetry and fit again
        double cn3rra_params[5], cn3raa_params[5], *cn3aa_params = cn2aa_params;
        double cn3aa_fit = cn2aa_fit, cn3ra_fit;
        char *cn3aa_func = cn2aa_func, *cn3ra_func;
        double min_dx3   = 0.5 - 1./(args->min_fraction+2);
        peakfit_reset(pkf);
        peakfit_add_bounded_gaussian(pkf, 1.0,1/3.,0.03, xrr,xra-min_dx3, 7);
        peakfit_set_mc(pkf, xrr,xra-min_dx3, 1,nmc);
        peakfit_add_bounded_gaussian(pkf, 1.0,2/3.,0.03, xra+min_dx3,xaa, 7);
        peakfit_set_mc(pkf, xra+min_dx3,xaa, 1,nmc);
        peakfit_run(pkf, nrr_aa, xrr_vals, yrr_vals);
        // force symmetry around x=0.5
        peakfit_get_params(pkf,0,cn3rra_params,5);
        peakfit_get_params(pkf,1,cn3raa_params,5);
        double cn3_dx = (0.5-cn3rra_params[1] + cn3raa_params[1]-0.5)*0.5;
        if ( cn3_dx > 0.5/3 ) cn3_dx = 0.5/3;   // CN3 peaks should not be separated by more than 1/3
        peakfit_reset(pkf);
        peakfit_add_gaussian(pkf, cn3rra_params[0],0.5-cn3_dx,cn3rra_params[2], 5);
        peakfit_add_gaussian(pkf, cn3raa_params[0],0.5+cn3_dx,cn3raa_params[2], 5);
        cn3ra_fit  = peakfit_run(pkf, nrr_aa, xrr_vals, yrr_vals);
        cn3ra_func = strdup(peakfit_sprint_func(pkf));
        // compare peak sizes
        peakfit_get_params(pkf,0,cn3rra_params,3);
        peakfit_get_params(pkf,1,cn3raa_params,3);
        double cn3rra_size = cn3rra_params[0]*cn3rra_params[0];
        double cn3raa_size = cn3raa_params[0]*cn3raa_params[0];
        double cn3_dy      = cn3rra_size > cn3raa_size ? cn3raa_size/cn3rra_size : cn3rra_size/cn3raa_size;
        double cn3_frac    = (1 - 2*cn3rra_params[1]) / cn3rra_params[1];
        double cn3_fit     = cn3ra_fit + cn3aa_fit;
        // A very reasonable heuristics: check if the peak's width converged, exclude far too broad or far too narrow peaks
        if ( cn3rra_params[2]>0.3  || cn3raa_params[2]>0.3 ) cn3_fit = HUGE_VAL;
        if ( cn3rra_params[2]<1e-2 || cn3raa_params[2]<1e-2 ) cn3_fit = HUGE_VAL;

        // CN4 (contaminations)
        // - first fit only the [0,0.5] part of the data, then enforce the symmetry and fit again
        // - min_frac=1 (resp. 0.5) is interpreted as 50:50% (rep. 75:25%) contamination
        double cn4AAaa_params[3] = {1,1,1} ,cn4AAra_params[3] = {1,1,1}, cn4RAra_params[3], cn4RArr_params[5], cn4RAaa_params[5];
        double cn4aa_fit = 0, cn4ra_fit;
        char *cn4aa_func = 0, *cn4ra_func;
        double min_dx4   = 0.25*args->min_fraction;
        if ( args->include_aa )
        {
            peakfit_reset(pkf);
            peakfit_add_exp(pkf, 0.5,1.0,0.2, 5);
            peakfit_set_mc(pkf, 0.01,0.3,2,nmc);
            peakfit_add_bounded_gaussian(pkf, 0.4,(xaa+xmax)*0.5,2e-2, xaa,xmax, 7);
            peakfit_set_mc(pkf, xaa,xmax, 1,nmc);
            cn4aa_fit  = peakfit_run(pkf, naa_max, xaa_vals,yaa_vals);
            cn4aa_func = strdup(peakfit_sprint_func(pkf));
            peakfit_get_params(pkf,0,cn4AAaa_params,3);
            peakfit_get_params(pkf,1,cn4AAra_params,5);
        }
        peakfit_reset(pkf);
        // first fit only the [0,0.5] part of the data
        peakfit_add_gaussian(pkf, 1.0,0.5,0.03, 5);
        peakfit_add_bounded_gaussian(pkf, 0.6,0.3,0.03, xrr,xra-min_dx4, 7);
        peakfit_set_mc(pkf, xrr,xra-min_dx4,2,nmc);
        peakfit_run(pkf, nrr_ra , xrr_vals, yrr_vals);
        // now forcet symmetry around x=0.5
        peakfit_get_params(pkf,0,cn4RAra_params,3);
        peakfit_get_params(pkf,1,cn4RArr_params,5);
        double cn4_dx = 0.5-cn4RArr_params[1];
        if ( cn4_dx > 0.25 ) cn4_dx = 0.25;   // CN4 peaks should not be separated by more than 0.5
        peakfit_reset(pkf);
        peakfit_add_gaussian(pkf, cn4RAra_params[0],0.5,cn4RAra_params[2], 5);
        peakfit_add_gaussian(pkf, cn4RArr_params[0],0.5-cn4_dx,cn4RArr_params[2], 5);
        peakfit_add_gaussian(pkf, cn4RArr_params[0],0.5+cn4_dx,cn4RArr_params[2], 5);
        peakfit_set_mc(pkf, 0.1,cn4RAra_params[0],0,nmc);
        peakfit_set_mc(pkf, 0.01,0.1,2,nmc);
        cn4ra_fit  = peakfit_run(pkf, nrr_aa , xrr_vals, yrr_vals);
        cn4ra_func = strdup(peakfit_sprint_func(pkf));
        peakfit_get_params(pkf,0,cn4RAra_params,3);
        peakfit_get_params(pkf,1,cn4RArr_params,3);
        peakfit_get_params(pkf,2,cn4RAaa_params,3);
        double cn4RAra_size = cn4RAra_params[0]==0 ? HUGE_VAL : cn4RAra_params[0]*cn4RAra_params[0];
        double cn4RArr_size = cn4RArr_params[0]*cn4RArr_params[0];
        double cn4RAaa_size = cn4RAaa_params[0]*cn4RAaa_params[0];
        double cn4RArr_dy   = cn4RArr_size < cn4RAra_size ? cn4RArr_size/cn4RAra_size : cn4RAra_size/cn4RArr_size;
        double cn4RAaa_dy   = cn4RAaa_size < cn4RAra_size ? cn4RAaa_size/cn4RAra_size : cn4RAra_size/cn4RAaa_size;
        double cn4_dy       = cn4RArr_dy < cn4RAaa_dy ? cn4RArr_dy/cn4RAaa_dy : cn4RAaa_dy/cn4RArr_dy;
        double cn4_ymin     = cn4RArr_size < cn4RAaa_size ? cn4RArr_size/cn4RAra_size : cn4RAaa_size/cn4RAra_size;
        cn4_dx              = (cn4RAaa_params[1]-0.5) - (0.5-cn4RArr_params[1]);
        double cn4_frac     = cn4RAaa_params[1] - cn4RArr_params[1];
        double cn4_fit      = cn4ra_fit + cn4aa_fit;
        // A very reasonable heuristics: check if the peak's width converged, exclude far too broad or far too narrow peaks
        if ( cn4RAra_params[2]>0.3 || cn4RArr_params[2]>0.3 || cn4RAaa_params[2]>0.3 ) cn4_fit = HUGE_VAL;
        if ( cn4RAra_params[2]<1e-2 || cn4RArr_params[2]<1e-2 || cn4RAaa_params[2]<1e-2 ) cn4_fit = HUGE_VAL;

        // Choose the best match
        char cn2_fail = '*', cn3_fail = '*', cn4_fail = '*';
        if ( cn2_fit > args->fit_th ) cn2_fail = 'f';

        if ( cn3_fit > args->fit_th ) cn3_fail = 'f';
        else if ( cn3_dy < args->peak_symmetry ) cn3_fail = 'y';    // size difference is too big

        if ( cn4_fit > args->fit_th ) cn4_fail = 'f';
        else if ( cn4_ymin < args->min_peak_size ) cn4_fail = 'y';      // side peak is too small
        else if ( cn4_dy < args->peak_symmetry ) cn4_fail = 'Y';    // size difference is too big
        else if ( cn4_dx > 0.1 ) cn4_fail = 'x';                    // side peaks placed assymetrically

        double cn = -1, fit = cn2_fit;
        if ( cn2_fail == '*' ) { cn = 2; fit = cn2_fit; }
        if ( cn3_fail == '*' )
        {
            // Use cn_penalty as a tiebreaker. If set to 0.3, cn3_fit must be 30% smaller than cn2_fit.
            if ( cn<0 || cn3_fit < (1-args->cn_penalty) * fit )
            {
                cn = 2 + cn3_frac;
                fit = cn3_fit;
                if ( cn2_fail=='*' ) cn2_fail = 'p';
            }
            else cn3_fail = 'p';
        }
        if ( cn4_fail == '*' )
        {
            if ( cn<0 || cn4_fit < (1-args->cn_penalty) * fit )
            {
                cn = 3 + cn4_frac;
                fit = cn4_fit;
                if ( cn2_fail=='*' ) cn2_fail = 'p';
                if ( cn3_fail=='*' ) cn3_fail = 'p';
            }
            else cn4_fail = 'p';
        }

        if ( args->verbose )
        {
            fprintf(stderr,"\tcn2 %c fit=%e\n", cn2_fail, cn2_fit);
            fprintf(stderr,"\t       .. %e\n", cn2ra_fit);
            fprintf(stderr,"\t            RA:   %f %f %f\n", cn2ra_params[0],cn2ra_params[1],cn2ra_params[2]);
            fprintf(stderr,"\t       .. %e\n", cn2aa_fit);
            fprintf(stderr,"\t            AA:   %f %f %f\n", cn2aa_params[0],cn2aa_params[1],cn2aa_params[2]);
            fprintf(stderr,"\t      func:\n");
            fprintf(stderr,"\t            %s\n", cn2ra_func);
            fprintf(stderr,"\t            %s\n", cn2aa_func);
            fprintf(stderr,"\n");
            fprintf(stderr,"\tcn3 %c fit=%e  frac=%f  symmetry=%f\n", cn3_fail, cn3_fit, cn3_frac, cn3_dy);
            fprintf(stderr,"\t       .. %e\n", cn3ra_fit);
            fprintf(stderr,"\t            RRA:  %f %f %f\n", cn3rra_params[0],cn3rra_params[1],cn3rra_params[2]);
            fprintf(stderr,"\t            RAA:  %f %f %f\n", cn3raa_params[0],cn3raa_params[1],cn3raa_params[2]);
            fprintf(stderr,"\t       .. %e\n", cn3aa_fit);
            fprintf(stderr,"\t            AAA:  %f %f %f\n", cn3aa_params[0],cn3aa_params[1],cn3aa_params[2]);
            fprintf(stderr,"\t      func:\n");
            fprintf(stderr,"\t            %s\n", cn3ra_func);
            fprintf(stderr,"\t            %s\n", cn3aa_func);
            fprintf(stderr,"\n");
            fprintf(stderr,"\tcn4 %c fit=%e  frac=%f  symmetry=%f ymin=%f\n", cn4_fail, cn4_fit, cn4_frac, cn4_dy, cn4_ymin);
            fprintf(stderr,"\t       .. %e\n", cn4ra_fit);
            fprintf(stderr,"\t            RArr:  %f %f %f\n", cn4RArr_params[0],cn4RArr_params[1],cn4RArr_params[2]);
            fprintf(stderr,"\t            RAra:  %f %f %f\n", cn4RAra_params[0],cn4RAra_params[1],cn4RAra_params[2]);
            fprintf(stderr,"\t            RAaa:  %f %f %f\n", cn4RAaa_params[0],cn4RAaa_params[1],cn4RAaa_params[2]);
            fprintf(stderr,"\t       .. %e\n", cn4aa_fit);
            fprintf(stderr,"\t            AAaa:  %f %f %f\n", cn4AAaa_params[0],cn4AAaa_params[1],cn4AAaa_params[2]);
            fprintf(stderr,"\t      func:\n");
            fprintf(stderr,"\t            %s\n", cn4ra_func);
            fprintf(stderr,"\t            %s\n", cn4aa_func);
            fprintf(stderr,"\n");
        }

        if ( args->force_cn==2 || cn2_fail == '*' )
        {
            fprintf(args->dat_fp,"FIT\t%s\t%e\t%d\t%d\t%s\n", dist->chr,cn2ra_fit,dist->irr,dist->iaa,cn2ra_func);
            if ( cn2aa_func ) fprintf(args->dat_fp,"FIT\t%s\t%e\t%d\t%d\t%s\n", dist->chr,cn2aa_fit,dist->iaa,dist->nvals-1,cn2aa_func);
        }
        if ( args->force_cn==3 || cn3_fail == '*' )
        {
            fprintf(args->dat_fp,"FIT\t%s\t%e\t%d\t%d\t%s\n", dist->chr,cn3ra_fit,dist->irr,dist->iaa,cn3ra_func);
            if ( cn3aa_func ) fprintf(args->dat_fp,"FIT\t%s\t%e\t%d\t%d\t%s\n", dist->chr,cn3aa_fit,dist->iaa,dist->nvals-1,cn3aa_func);
        }
        if ( args->force_cn==4 || cn4_fail == '*' )
        {
            fprintf(args->dat_fp,"FIT\t%s\t%e\t%d\t%d\t%s\n", dist->chr,cn4ra_fit,dist->irr,dist->iaa,cn4ra_func);
            if ( cn4aa_func ) fprintf(args->dat_fp,"FIT\t%s\t%e\t%d\t%d\t%s\n", dist->chr,cn4aa_fit,dist->iaa,dist->nvals-1,cn4aa_func);
        }
        fprintf(args->dat_fp,"CN\t%s\t%.2f\t%f\n", dist->chr, cn, fit);

        free(cn2aa_func);
        free(cn2ra_func);
        free(cn3ra_func);
        free(cn4ra_func);
        free(cn4aa_func);
    }

    peakfit_destroy(pkf);
}

static void usage(args_t *args)
{
    fprintf(stderr, "\n");
    fprintf(stderr, "About:   Detect number of chromosomal copies from Illumina's B-allele frequency (BAF)\n");
    fprintf(stderr, "Usage:   bcftools polysomy [OPTIONS] <file.vcf>\n");
    fprintf(stderr, "\n");
    fprintf(stderr, "General options:\n");
    fprintf(stderr, "    -o, --output-dir <path>        \n");
    fprintf(stderr, "    -r, --regions <region>         restrict to comma-separated list of regions\n");
    fprintf(stderr, "    -R, --regions-file <file>      restrict to regions listed in a file\n");
    fprintf(stderr, "    -s, --sample <name>            sample to analyze\n");
    fprintf(stderr, "    -t, --targets <region>         similar to -r but streams rather than index-jumps\n");
    fprintf(stderr, "    -T, --targets-file <file>      similar to -R but streams rather than index-jumps\n");
    fprintf(stderr, "    -v, --verbose                  \n");
    fprintf(stderr, "\n");
    fprintf(stderr, "Algorithm options:\n");
    fprintf(stderr, "    -b, --peak-size <float>        minimum peak size (0-1, larger is stricter) [0.1]\n");
    fprintf(stderr, "    -c, --cn-penalty <float>       penalty for increasing CN (0-1, larger is stricter) [0.7]\n");
    fprintf(stderr, "    -f, --fit-th <float>           goodness of fit threshold (>0, smaller is stricter) [3.3]\n");
    fprintf(stderr, "    -i, --include-aa               include the AA peak in CN2 and CN3 evaluation\n");
    fprintf(stderr, "    -m, --min-fraction <float>     minimum distinguishable fraction of aberrant cells [0.1]\n");
    fprintf(stderr, "    -p, --peak-symmetry <float>    peak symmetry threshold (0-1, larger is stricter) [0.5]\n");
    fprintf(stderr, "\n");
    exit(1);
}

int main_polysomy(int argc, char *argv[])
{
    args_t *args = (args_t*) calloc(1,sizeof(args_t));
    args->argc   = argc; args->argv = argv;
    args->nbins  = 150;
    args->fit_th = 3.3;
    args->cn_penalty = 0.7;
    args->peak_symmetry = 0.5;
    args->min_peak_size = 0.1;
    args->ra_rr_scaling = 1;
    args->min_fraction = 0.1;
    args->smooth = -3;

    static struct option loptions[] =
    {
        {"ra-rr-scaling",0,0,1},    // hidden option
        {"force-cn",1,0,2},         // hidden option
        {"smooth",1,0,'S'},         // hidden option
        {"nbins",1,0,'n'},          // hidden option
        {"include-aa",0,0,'i'},
        {"peak-size",1,0,'b'},
        {"min-fraction",1,0,'m'},
        {"verbose",0,0,'v'},
        {"fit-th",1,0,'f'},
        {"cn-penalty",1,0,'c'},
        {"peak-symmetry",1,0,'p'},
        {"output-dir",1,0,'o'},
        {"sample",1,0,'s'},
        {"targets",1,0,'t'},
        {"targets-file",1,0,'T'},
        {"regions",1,0,'r'},
        {"regions-file",1,0,'R'},
        {0,0,0,0}
    };
    char c, *tmp;
    while ((c = getopt_long(argc, argv, "h?o:vt:T:r:R:s:f:p:c:im:b:n:S:",loptions,NULL)) >= 0)
    {
        switch (c)
        {
            case  1 : args->ra_rr_scaling = 0; break;
            case  2 : args->force_cn = atoi(optarg); break;
            case 'n': args->nbins = atoi(optarg); break;
            case 'S': args->smooth = atoi(optarg); break;
            case 'i': args->include_aa = 1; break;
            case 'b':
                args->min_peak_size = strtod(optarg,&tmp);
                if ( *tmp ) error("Could not parse: -b %s\n", optarg);
                if ( args->min_peak_size<0 || args->min_peak_size>1 ) error("Range error: -b %s\n", optarg);
                break;
            case 'm':
                args->min_fraction = strtod(optarg,&tmp);
                if ( *tmp ) error("Could not parse: -n %s\n", optarg);
                if ( args->min_fraction<0 || args->min_fraction>1 ) error("Range error: -n %s\n", optarg);
                break;
            case 'f':
                args->fit_th = strtod(optarg,&tmp);
                if ( *tmp ) error("Could not parse: -f %s\n", optarg);
                break;
            case 'p':
                args->peak_symmetry = strtod(optarg,&tmp);
                if ( *tmp ) error("Could not parse: -p %s\n", optarg);
                break;
            case 'c':
                args->cn_penalty = strtod(optarg,&tmp);
                if ( *tmp ) error("Could not parse: -c %s\n", optarg);
                break;
            case 's': args->sample = optarg; break;
            case 't': args->targets_list = optarg; break;
            case 'T': args->targets_list = optarg; args->targets_is_file = 1; break;
            case 'r': args->regions_list = optarg; break;
            case 'R': args->regions_list = optarg; args->regions_is_file = 1; break;
            case 'o': args->output_dir = optarg; break;
            case 'v': args->verbose++; break;
            default: usage(args); break;
        }
    }
    if ( optind>=argc )
    {
        if ( !isatty(fileno((FILE *)stdin)) ) args->fname = "-";
    }
    else args->fname = argv[optind];
    if ( !args->fname ) usage(args);
    if ( !args->output_dir ) error("Missing the -o option\n");

    init_data(args);
    fit_curves(args);
    destroy_data(args);
    free(args);

    return 0;
}


/* multilarge_nlinear/fdf.c
 *
 * Copyright (C) 2015, 2016 Patrick Alken
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <string.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multilarge_nlinear.h>

gsl_multilarge_nlinear_workspace *
gsl_multilarge_nlinear_alloc (const gsl_multilarge_nlinear_type * T,
                              const gsl_multilarge_nlinear_parameters * params,
                              const size_t n, const size_t p)
{
  gsl_multilarge_nlinear_workspace * w;

  if (n < p)
    {
      GSL_ERROR_VAL ("insufficient data points, n < p", GSL_EINVAL, 0);
    }

  w = calloc (1, sizeof (gsl_multilarge_nlinear_workspace));
  if (w == 0)
    {
      GSL_ERROR_VAL ("failed to allocate space for workspace",
                     GSL_ENOMEM, 0);
    }

  w->n = n;
  w->p = p;
  w->type = T;
  w->fdf = NULL;
  w->niter = 0;
  w->params = *params;

  /* the cgst method uses its own built-in linear solver */
  if (w->params.trs == gsl_multilarge_nlinear_trs_cgst)
    {
      w->params.solver = gsl_multilarge_nlinear_solver_none;
    }

  w->x = gsl_vector_calloc (p);
  if (w->x == 0)
    {
      gsl_multilarge_nlinear_free (w);
      GSL_ERROR_VAL ("failed to allocate space for x", GSL_ENOMEM, 0);
    }

  w->f = gsl_vector_calloc (n);
  if (w->f == 0)
    {
      gsl_multilarge_nlinear_free (w);
      GSL_ERROR_VAL ("failed to allocate space for f", GSL_ENOMEM, 0);
    }

  w->dx = gsl_vector_calloc (p);
  if (w->dx == 0)
    {
      gsl_multilarge_nlinear_free (w);
      GSL_ERROR_VAL ("failed to allocate space for dx", GSL_ENOMEM, 0);
    }

  w->g = gsl_vector_alloc (p);
  if (w->g == 0)
    {
      gsl_multilarge_nlinear_free (w);
      GSL_ERROR_VAL ("failed to allocate space for g", GSL_ENOMEM, 0);
    }

  if (w->params.solver == gsl_multilarge_nlinear_solver_cholesky ||
      w->params.solver == gsl_multilarge_nlinear_solver_mcholesky)
    {
      w->JTJ = gsl_matrix_alloc (p, p);
      if (w->JTJ == 0)
        {
          gsl_multilarge_nlinear_free (w);
          GSL_ERROR_VAL ("failed to allocate space for JTJ", GSL_ENOMEM, 0);
        }
    }

  w->sqrt_wts_work = gsl_vector_calloc (n);
  if (w->sqrt_wts_work == 0)
    {
      gsl_multilarge_nlinear_free (w);
      GSL_ERROR_VAL ("failed to allocate space for weights", GSL_ENOMEM, 0);
    }

  w->state = (T->alloc)(&(w->params), n, p);
  if (w->state == 0)
    {
      gsl_multilarge_nlinear_free (w);
      GSL_ERROR_VAL ("failed to allocate space for state", GSL_ENOMEM, 0);
    }

  return w;
}

void
gsl_multilarge_nlinear_free (gsl_multilarge_nlinear_workspace * w)
{
  RETURN_IF_NULL (w);

  if (w->state)
    (w->type->free) (w->state);

  if (w->dx)
    gsl_vector_free (w->dx);

  if (w->x)
    gsl_vector_free (w->x);

  if (w->f)
    gsl_vector_free (w->f);

  if (w->sqrt_wts_work)
    gsl_vector_free (w->sqrt_wts_work);

  if (w->g)
    gsl_vector_free (w->g);

  if (w->JTJ)
    gsl_matrix_free (w->JTJ);

  free (w);
}

gsl_multilarge_nlinear_parameters
gsl_multilarge_nlinear_default_parameters(void)
{
  gsl_multilarge_nlinear_parameters params;

  params.trs = gsl_multilarge_nlinear_trs_lm;
  params.scale = gsl_multilarge_nlinear_scale_more;
  params.solver = gsl_multilarge_nlinear_solver_cholesky;
  params.fdtype = GSL_MULTILARGE_NLINEAR_FWDIFF;
  params.factor_up = 3.0;
  params.factor_down = 2.0;
  params.avmax = 0.75;
  params.h_df = GSL_SQRT_DBL_EPSILON;
  params.h_fvv = 0.01;
  params.max_iter = 0;
  params.tol = 1.0e-6;

  return params;
}

int
gsl_multilarge_nlinear_init (const gsl_vector * x,
                             gsl_multilarge_nlinear_fdf * fdf,
                             gsl_multilarge_nlinear_workspace * w)
{
  return gsl_multilarge_nlinear_winit(x, NULL, fdf, w);
}

int
gsl_multilarge_nlinear_winit (const gsl_vector * x,
                              const gsl_vector * wts,
                              gsl_multilarge_nlinear_fdf * fdf,
                              gsl_multilarge_nlinear_workspace * w)
{
  const size_t n = w->f->size;

  if (n != fdf->n)
    {
      GSL_ERROR ("function size does not match workspace", GSL_EBADLEN);
    }
  else if (w->x->size != x->size)
    {
      GSL_ERROR ("vector length does not match workspace", GSL_EBADLEN);
    }
  else if (wts != NULL && n != wts->size)
    {
      GSL_ERROR ("weight vector length does not match workspace", GSL_EBADLEN);
    }
  else
    {
      size_t i;

      /* initialize counters for function and Jacobian evaluations */
      fdf->nevalf = 0;
      fdf->nevaldfu = 0;
      fdf->nevaldf2 = 0;
      fdf->nevalfvv = 0;

      w->fdf = fdf;
      gsl_vector_memcpy(w->x, x);
      w->niter = 0;

      if (wts)
        {
          w->sqrt_wts = w->sqrt_wts_work;

          for (i = 0; i < n; ++i)
            {
              double wi = gsl_vector_get(wts, i);
              gsl_vector_set(w->sqrt_wts, i, sqrt(wi));
            }
        }
      else
        {
          w->sqrt_wts = NULL;
        }

      return (w->type->init) (w->state, w->sqrt_wts, w->fdf,
                              w->x, w->f, w->g, w->JTJ);
    }
}

int
gsl_multilarge_nlinear_iterate (gsl_multilarge_nlinear_workspace * w)
{
  int status =
    (w->type->iterate) (w->state, w->sqrt_wts, w->fdf,
                        w->x, w->f, w->g, w->JTJ, w->dx);

  w->niter++;

  return status;
}

double
gsl_multilarge_nlinear_avratio (const gsl_multilarge_nlinear_workspace * w)
{
  return (w->type->avratio) (w->state);
}

int
gsl_multilarge_nlinear_rcond (double * rcond, const gsl_multilarge_nlinear_workspace * w)
{
  int status = (w->type->rcond) (rcond, w->JTJ, w->state);
  return status;
}

int
gsl_multilarge_nlinear_covar (gsl_matrix * covar, gsl_multilarge_nlinear_workspace * w)
{
  if (covar->size1 != covar->size2)
    {
      GSL_ERROR ("covariance matrix must be square", GSL_ENOTSQR);
    }
  else if (covar->size1 != w->p)
    {
      GSL_ERROR ("covariance matrix does not match workspace", GSL_EBADLEN);
    }
  else
    {
      int status = (w->type->covar) (w->JTJ, covar, w->state);
      return status;
    }
}

/*
gsl_multilarge_nlinear_driver()
  Iterate the nonlinear least squares solver until completion

Inputs: maxiter  - maximum iterations to allow
        xtol     - tolerance in step x
        gtol     - tolerance in gradient
        ftol     - tolerance in ||f||
        callback - callback function to call each iteration
        callback_params - parameters to pass to callback function
        info     - (output) info flag on why iteration terminated
                   1 = stopped due to small step size ||dx|
                   2 = stopped due to small gradient
                   3 = stopped due to small change in f
                   GSL_ETOLX = ||dx|| has converged to within machine
                               precision (and xtol is too small)
                   GSL_ETOLG = ||g||_inf is smaller than machine
                               precision (gtol is too small)
                   GSL_ETOLF = change in ||f|| is smaller than machine
                               precision (ftol is too small)
        w        - workspace

Return:
GSL_SUCCESS if converged
GSL_MAXITER if maxiter exceeded without converging
GSL_ENOPROG if no accepted step found on first iteration
*/

int
gsl_multilarge_nlinear_driver (const size_t maxiter,
                               const double xtol,
                               const double gtol,
                               const double ftol,
                               void (*callback)(const size_t iter, void *params,
                                                const gsl_multilarge_nlinear_workspace *w),
                               void *callback_params,
                               int *info,
                               gsl_multilarge_nlinear_workspace * w)
{
  int status;
  size_t iter = 0;

  /* call user callback function prior to any iterations
   * with initial system state */
  if (callback)
    callback(iter, callback_params, w);

  do
    {
      status = gsl_multilarge_nlinear_iterate (w);

      /*
       * If the solver reports no progress on the first iteration,
       * then it didn't find a single step to reduce the
       * cost function and more iterations won't help so return.
       *
       * If we get a no progress flag on subsequent iterations,
       * it means we did find a good step in a previous iteration,
       * so continue iterating since the solver has now reset
       * mu to its initial value.
       */
      if (status == GSL_ENOPROG && iter == 0)
        {
          *info = status;
          return GSL_EMAXITER;
        }

      ++iter;

      if (callback)
        callback(iter, callback_params, w);

      /* test for convergence */
      status = gsl_multilarge_nlinear_test(xtol, gtol, ftol, info, w);
    }
  while (status == GSL_CONTINUE && iter < maxiter);

  /*
   * the following error codes mean that the solution has converged
   * to within machine precision, so record the error code in info
   * and return success
   */
  if (status == GSL_ETOLF || status == GSL_ETOLX || status == GSL_ETOLG)
    {
      *info = status;
      status = GSL_SUCCESS;
    }

  /* check if max iterations reached */
  if (iter >= maxiter && status != GSL_SUCCESS)
    status = GSL_EMAXITER;

  return status;
} /* gsl_multilarge_nlinear_driver() */

const char *
gsl_multilarge_nlinear_name (const gsl_multilarge_nlinear_workspace * w)
{
  return w->type->name;
}

gsl_vector *
gsl_multilarge_nlinear_position (const gsl_multilarge_nlinear_workspace * w)
{
  return w->x;
}

gsl_vector *
gsl_multilarge_nlinear_residual (const gsl_multilarge_nlinear_workspace * w)
{
  return w->f;
}

gsl_vector *
gsl_multilarge_nlinear_step (const gsl_multilarge_nlinear_workspace * w)
{
  return w->dx;
}

size_t
gsl_multilarge_nlinear_niter (const gsl_multilarge_nlinear_workspace * w)
{
  return w->niter;
}

const char *
gsl_multilarge_nlinear_trs_name (const gsl_multilarge_nlinear_workspace * w)
{
  return w->params.trs->name;
}

/*
gsl_multilarge_nlinear_eval_f()
  Compute residual vector y with user callback function, and apply
weighting transform if given:

y~ = sqrt(W) y

Inputs: fdf  - callback function
        x    - model parameters
        swts - weight matrix sqrt(W) = sqrt(diag(w1,w2,...,wn))
               set to NULL for unweighted fit
        y    - (output) (weighted) residual vector
               y_i = sqrt(w_i) f_i where f_i is unweighted residual
*/

int
gsl_multilarge_nlinear_eval_f(gsl_multilarge_nlinear_fdf *fdf,
                              const gsl_vector *x,
                              const gsl_vector *swts,
                              gsl_vector *y)
{
  int s = ((*((fdf)->f)) (x, fdf->params, y));

  ++(fdf->nevalf);

  /* y <- sqrt(W) y */
  if (swts)
    gsl_vector_mul(y, swts);

  return s;
}

/*
gsl_multilarge_nlinear_eval_df()
  Compute Jacobian matrix-vector product:

v = J * u

or

v = J^T u

Inputs: TransJ - use J or J^T
        x      - model parameters
        f      - residual vector f(x)
        u      - input vector u
        swts   - weight matrix W = diag(w1,w2,...,wn)
                 set to NULL for unweighted fit
        h      - finite difference step size
        fdtype - finite difference method
        fdf    - callback function
        v      - (output) vector v
        JTJ    - (output) matrix J^T J
        work   - workspace for finite difference, size n
*/

int
gsl_multilarge_nlinear_eval_df(const CBLAS_TRANSPOSE_t TransJ,
                               const gsl_vector *x,
                               const gsl_vector *f,
                               const gsl_vector *u,
                               const gsl_vector *swts,
                               const double h,
                               const gsl_multilarge_nlinear_fdtype fdtype,
                               gsl_multilarge_nlinear_fdf *fdf,
                               gsl_vector *v,
                               gsl_matrix *JTJ,
                               gsl_vector *work)
{
  const size_t n = fdf->n;
  const size_t p = fdf->p;

  if (u != NULL && ((TransJ == CblasNoTrans && u->size != p) ||
                    (TransJ == CblasTrans && u->size != n)))
    {
      GSL_ERROR("u vector has wrong size", GSL_EBADLEN);
    }
  else if (v != NULL && ((TransJ == CblasNoTrans && v->size != n) ||
                         (TransJ == CblasTrans && v->size != p)))
    {
      GSL_ERROR("v vector has wrong size", GSL_EBADLEN);
    }
  else if (JTJ != NULL && ((JTJ->size1 != p) || (JTJ->size2 != p)))
    {
      GSL_ERROR("JTJ matrix has wrong size", GSL_EBADLEN);
    }
  else
    {
      int status = GSL_SUCCESS;

      if (fdf->df)
        {
          /* call user-supplied function */
          status = ((*((fdf)->df)) (TransJ, x, u, fdf->params, v, JTJ));

          if (v)
            ++(fdf->nevaldfu);

          if (JTJ)
            ++(fdf->nevaldf2);
        }
      else
        {
#if 0
          /* use finite difference Jacobian approximation */
          status = gsl_multilarge_nlinear_df(h, fdtype, x, swts, fdf, f, df, work);
#endif
        }

      return status;
    }
}

/*
gsl_multilarge_nlinear_eval_fvv()
  Compute second direction derivative vector yvv with user
callback function, and apply weighting transform if given:

yvv~ = sqrt(W) yvv

Inputs: h    - step size for finite difference, if needed
        x    - model parameters, size p
        v    - unscaled geodesic velocity vector, size p
        f    - residual vector f(x), size n
        swts - weight matrix sqrt(W) = sqrt(diag(w1,w2,...,wn))
               set to NULL for unweighted fit
        fdf  - callback function
        yvv  - (output) (weighted) second directional derivative vector
               yvv_i = sqrt(w_i) fvv_i where f_i is unweighted
        work - workspace, size p
*/

int
gsl_multilarge_nlinear_eval_fvv(const double h,
                                const gsl_vector *x,
                                const gsl_vector *v,
                                const gsl_vector *f,
                                const gsl_vector *swts,
                                gsl_multilarge_nlinear_fdf *fdf,
                                gsl_vector *yvv,
                                gsl_vector *work)
{
  int status = GSL_SUCCESS;

  if (fdf->fvv != NULL)
    {
      /* call user-supplied function */
      status = ((*((fdf)->fvv)) (x, v, fdf->params, yvv));
      ++(fdf->nevalfvv);
    }
  else
    {
#if 0
      /* use finite difference approximation */
      status = gsl_multilarge_nlinear_fdfvv(h, x, v, f, J,
                                          swts, fdf, yvv, work);
#endif
    }

  /* yvv <- sqrt(W) yvv */
  if (swts)
    gsl_vector_mul(yvv, swts);

  return status;
}


// Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
// the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
// reserved. See files LICENSE and NOTICE for details.
//
// This file is part of CEED, a collection of benchmarks, miniapps, software
// libraries and APIs for efficient high-order finite element and spectral
// element discretizations for exascale applications. For more information and
// source code availability see http://github.com/ceed.
//
// The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
// a collaborative effort of two U.S. Department of Energy organizations (Office
// of Science and the National Nuclear Security Administration) responsible for
// the planning and preparation of a capable exascale ecosystem, including
// software, applications, hardware, advanced system engineering and early
// testbed platforms, in support of the nation's exascale computing imperative.

#ifndef setupsphere_h
#define setupsphere_h

#include <ceed.h>
#include <petsc.h>
#include <petscdmplex.h>
#include <petscfe.h>
#include <stdbool.h>
#include <string.h>
#include "qfunctions/bps/bp1sphere.h"
#include "qfunctions/bps/bp2sphere.h"
#include "qfunctions/bps/bp3sphere.h"
#include "qfunctions/bps/bp4sphere.h"
#include "qfunctions/bps/common.h"

#if PETSC_VERSION_LT(3,14,0)
#  define DMPlexGetClosureIndices(a,b,c,d,e,f,g,h,i) DMPlexGetClosureIndices(a,b,c,d,f,g,i)
#  define DMPlexRestoreClosureIndices(a,b,c,d,e,f,g,h,i) DMPlexRestoreClosureIndices(a,b,c,d,f,g,i)
#endif

#if PETSC_VERSION_LT(3,14,0)
#  define DMPlexCreateSphereMesh(a,b,c,d,e) DMPlexCreateSphereMesh(a,b,c,e)
#endif

// -----------------------------------------------------------------------------
// PETSc Operator Structs
// -----------------------------------------------------------------------------

// Data for PETSc
typedef struct UserO_ *UserO;
struct UserO_ {
  MPI_Comm comm;
  DM dm;
  Vec Xloc, Yloc, diag;
  CeedVector xceed, yceed;
  CeedOperator op;
  Ceed ceed;
};

// Data for PETSc Interp/Restrict operators
typedef struct UserIR_ *UserIR;
struct UserIR_ {
  MPI_Comm comm;
  DM dmc, dmf;
  Vec Xloc, Yloc, mult;
  CeedVector ceedvecc, ceedvecf;
  CeedOperator op;
  Ceed ceed;
};

// -----------------------------------------------------------------------------
// libCEED Data Struct
// -----------------------------------------------------------------------------

// libCEED data struct for level
typedef struct CeedData_ *CeedData;
struct CeedData_ {
  Ceed ceed;
  CeedBasis basisx, basisu, basisctof;
  CeedElemRestriction Erestrictx, Erestrictu, Erestrictui, Erestrictqdi;
  CeedQFunction qf_apply;
  CeedOperator op_apply, op_restrict, op_interp;
  CeedVector qdata, xceed, yceed;
};

// -----------------------------------------------------------------------------
// Command Line Options
// -----------------------------------------------------------------------------

// Coarsening options
typedef enum {
  COARSEN_UNIFORM = 0, COARSEN_LOGARITHMIC = 1
} coarsenType;
static const char *const coarsenTypes [] = {"uniform","logarithmic",
                                            "coarsenType","COARSEN",0
                                           };

// -----------------------------------------------------------------------------
// BP Option Data
// -----------------------------------------------------------------------------

// BP options
typedef enum {
  CEED_BP1 = 0, CEED_BP2 = 1, CEED_BP3 = 2,
  CEED_BP4 = 3, CEED_BP5 = 4, CEED_BP6 = 5
} bpType;
static const char *const bpTypes[] = {"bp1","bp2","bp3","bp4","bp5","bp6",
                                      "bpType","CEED_BP",0
                                     };

// BP specific data
typedef struct {
  CeedInt ncompu, qdatasize, qextra;
  CeedQFunctionUser setupgeo, setuprhs, apply, error;
  const char *setupgeofname, *setuprhsfname, *applyfname, *errorfname;
  CeedEvalMode inmode, outmode;
  CeedQuadMode qmode;
  PetscBool enforce_bc;
  PetscErrorCode (*bcs_func)(PetscInt, PetscReal, const PetscReal *,
                             PetscInt, PetscScalar *, void *);
} bpData;

static bpData bpOptions[6] = {
  [CEED_BP1] = {
    .ncompu = 1,
    .qdatasize = 1,
    .qextra = 1,
    .setupgeo = SetupMassGeo,
    .setuprhs = SetupMassRhs,
    .apply = Mass,
    .error = Error,
    .setupgeofname = SetupMassGeo_loc,
    .setuprhsfname = SetupMassRhs_loc,
    .applyfname = Mass_loc,
    .errorfname = Error_loc,
    .inmode = CEED_EVAL_INTERP,
    .outmode = CEED_EVAL_INTERP,
    .qmode = CEED_GAUSS
  },
  [CEED_BP2] = {
    .ncompu = 3,
    .qdatasize = 1,
    .qextra = 1,
    .setupgeo = SetupMassGeo,
    .setuprhs = SetupMassRhs3,
    .apply = Mass3,
    .error = Error3,
    .setupgeofname = SetupMassGeo_loc,
    .setuprhsfname = SetupMassRhs3_loc,
    .applyfname = Mass3_loc,
    .errorfname = Error3_loc,
    .inmode = CEED_EVAL_INTERP,
    .outmode = CEED_EVAL_INTERP,
    .qmode = CEED_GAUSS
  },
  [CEED_BP3] = {
    .ncompu = 1,
    .qdatasize = 4,
    .qextra = 1,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs,
    .apply = Diff,
    .error = Error,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs_loc,
    .applyfname = Diff_loc,
    .errorfname = Error_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS
  },
  [CEED_BP4] = {
    .ncompu = 3,
    .qdatasize = 4,
    .qextra = 1,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs3,
    .apply = Diff3,
    .error = Error3,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs3_loc,
    .applyfname = Diff_loc,
    .errorfname = Error3_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS
  },
  [CEED_BP5] = {
    .ncompu = 1,
    .qdatasize = 4,
    .qextra = 0,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs,
    .apply = Diff,
    .error = Error,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs_loc,
    .applyfname = Diff_loc,
    .errorfname = Error_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS_LOBATTO
  },
  [CEED_BP6] = {
    .ncompu = 3,
    .qdatasize = 4,
    .qextra = 0,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs3,
    .apply = Diff3,
    .error = Error3,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs3_loc,
    .applyfname = Diff_loc,
    .errorfname = Error3_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS_LOBATTO
  }
};

// -----------------------------------------------------------------------------
// PETSc sphere auxiliary functions
// -----------------------------------------------------------------------------

// Utility function taken from petsc/src/dm/impls/plex/examples/tutorials/ex7.c
static PetscErrorCode ProjectToUnitSphere(DM dm) {
  Vec            coordinates;
  PetscScalar   *coords;
  PetscInt       Nv, v, dim, d;
  PetscErrorCode ierr;

  PetscFunctionBeginUser;
  ierr = DMGetCoordinatesLocal(dm, &coordinates); CHKERRQ(ierr);
  ierr = VecGetLocalSize(coordinates, &Nv); CHKERRQ(ierr);
  ierr = VecGetBlockSize(coordinates, &dim); CHKERRQ(ierr);
  Nv  /= dim;
  ierr = VecGetArray(coordinates, &coords); CHKERRQ(ierr);
  for (v = 0; v < Nv; ++v) {
    PetscReal r = 0.0;

    for (d = 0; d < dim; ++d) r += PetscSqr(PetscRealPart(coords[v*dim+d]));
    r = PetscSqrtReal(r);
    for (d = 0; d < dim; ++d) coords[v*dim+d] /= r;
  }
  ierr = VecRestoreArray(coordinates, &coords); CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

// -----------------------------------------------------------------------------
// PETSc FE Boilerplate
// -----------------------------------------------------------------------------

// Create FE by degree
static int PetscFECreateByDegree(DM dm, PetscInt dim, PetscInt Nc,
                                 PetscBool isSimplex, const char prefix[],
                                 PetscInt order, PetscFE *fem) {
  PetscQuadrature q, fq;
  DM              K;
  PetscSpace      P;
  PetscDualSpace  Q;
  PetscInt        quadPointsPerEdge;
  PetscBool       tensor = isSimplex ? PETSC_FALSE : PETSC_TRUE;
  PetscErrorCode  ierr;

  PetscFunctionBeginUser;
  /* Create space */
  ierr = PetscSpaceCreate(PetscObjectComm((PetscObject) dm), &P); CHKERRQ(ierr);
  ierr = PetscObjectSetOptionsPrefix((PetscObject) P, prefix); CHKERRQ(ierr);
  ierr = PetscSpacePolynomialSetTensor(P, tensor); CHKERRQ(ierr);
  ierr = PetscSpaceSetFromOptions(P); CHKERRQ(ierr);
  ierr = PetscSpaceSetNumComponents(P, Nc); CHKERRQ(ierr);
  ierr = PetscSpaceSetNumVariables(P, dim); CHKERRQ(ierr);
  ierr = PetscSpaceSetDegree(P, order, order); CHKERRQ(ierr);
  ierr = PetscSpaceSetUp(P); CHKERRQ(ierr);
  ierr = PetscSpacePolynomialGetTensor(P, &tensor); CHKERRQ(ierr);
  /* Create dual space */
  ierr = PetscDualSpaceCreate(PetscObjectComm((PetscObject) dm), &Q);
  CHKERRQ(ierr);
  ierr = PetscDualSpaceSetType(Q,PETSCDUALSPACELAGRANGE); CHKERRQ(ierr);
  ierr = PetscObjectSetOptionsPrefix((PetscObject) Q, prefix); CHKERRQ(ierr);
  ierr = PetscDualSpaceCreateReferenceCell(Q, dim, isSimplex, &K); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetDM(Q, K); CHKERRQ(ierr);
  ierr = DMDestroy(&K); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetNumComponents(Q, Nc); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetOrder(Q, order); CHKERRQ(ierr);
  ierr = PetscDualSpaceLagrangeSetTensor(Q, tensor); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetFromOptions(Q); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetUp(Q); CHKERRQ(ierr);
  /* Create element */
  ierr = PetscFECreate(PetscObjectComm((PetscObject) dm), fem); CHKERRQ(ierr);
  ierr = PetscObjectSetOptionsPrefix((PetscObject) *fem, prefix); CHKERRQ(ierr);
  ierr = PetscFESetFromOptions(*fem); CHKERRQ(ierr);
  ierr = PetscFESetBasisSpace(*fem, P); CHKERRQ(ierr);
  ierr = PetscFESetDualSpace(*fem, Q); CHKERRQ(ierr);
  ierr = PetscFESetNumComponents(*fem, Nc); CHKERRQ(ierr);
  ierr = PetscFESetUp(*fem); CHKERRQ(ierr);
  ierr = PetscSpaceDestroy(&P); CHKERRQ(ierr);
  ierr = PetscDualSpaceDestroy(&Q); CHKERRQ(ierr);
  /* Create quadrature */
  quadPointsPerEdge = PetscMax(order + 1,1);
  if (isSimplex) {
    ierr = PetscDTStroudConicalQuadrature(dim,   1, quadPointsPerEdge, -1.0, 1.0,
                                          &q); CHKERRQ(ierr);
    ierr = PetscDTStroudConicalQuadrature(dim-1, 1, quadPointsPerEdge, -1.0, 1.0,
                                          &fq); CHKERRQ(ierr);
  } else {
    ierr = PetscDTGaussTensorQuadrature(dim,   1, quadPointsPerEdge, -1.0, 1.0,
                                        &q); CHKERRQ(ierr);
    ierr = PetscDTGaussTensorQuadrature(dim-1, 1, quadPointsPerEdge, -1.0, 1.0,
                                        &fq); CHKERRQ(ierr);
  }
  ierr = PetscFESetQuadrature(*fem, q); CHKERRQ(ierr);
  ierr = PetscFESetFaceQuadrature(*fem, fq); CHKERRQ(ierr);
  ierr = PetscQuadratureDestroy(&q); CHKERRQ(ierr);
  ierr = PetscQuadratureDestroy(&fq); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// -----------------------------------------------------------------------------
// PETSc Setup for Level
// -----------------------------------------------------------------------------

// This function sets up a DM for a given degree
static int SetupDMByDegree(DM dm, PetscInt degree, PetscInt ncompu,
                           PetscInt dim) {
  PetscInt ierr;
  PetscFE fe;

  PetscFunctionBeginUser;

  // Setup FE
  ierr = PetscFECreateByDegree(dm, dim, ncompu, PETSC_FALSE, NULL, degree, &fe);
  CHKERRQ(ierr);
  ierr = DMAddField(dm, NULL, (PetscObject)fe); CHKERRQ(ierr);

  // Setup DM
  ierr = DMCreateDS(dm); CHKERRQ(ierr);
  ierr = DMPlexSetClosurePermutationTensor(dm, PETSC_DETERMINE, NULL);
  CHKERRQ(ierr);
  ierr = PetscFEDestroy(&fe); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// -----------------------------------------------------------------------------
// libCEED Setup for Level
// -----------------------------------------------------------------------------

// Destroy libCEED operator objects
static PetscErrorCode CeedDataDestroy(CeedInt i, CeedData data) {
  PetscInt ierr;

  CeedVectorDestroy(&data->qdata);
  CeedVectorDestroy(&data->xceed);
  CeedVectorDestroy(&data->yceed);
  CeedBasisDestroy(&data->basisx);
  CeedBasisDestroy(&data->basisu);
  CeedElemRestrictionDestroy(&data->Erestrictu);
  CeedElemRestrictionDestroy(&data->Erestrictx);
  CeedElemRestrictionDestroy(&data->Erestrictui);
  CeedElemRestrictionDestroy(&data->Erestrictqdi);
  CeedQFunctionDestroy(&data->qf_apply);
  CeedOperatorDestroy(&data->op_apply);
  if (i > 0) {
    CeedOperatorDestroy(&data->op_interp);
    CeedBasisDestroy(&data->basisctof);
    CeedOperatorDestroy(&data->op_restrict);
  }
  ierr = PetscFree(data); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// Get CEED restriction data from DMPlex
static int CreateRestrictionPlex(Ceed ceed, CeedInt P, CeedInt ncomp,
                                 CeedElemRestriction *Erestrict, DM dm) {
  PetscInt ierr;
  PetscInt c, cStart, cEnd, nelem, nnodes, *erestrict, eoffset;
  PetscSection section;
  Vec Uloc;

  PetscFunctionBegin;

  // Get Nelem
  ierr = DMGetSection(dm, &section); CHKERRQ(ierr);
  ierr = DMPlexGetHeightStratum(dm, 0, &cStart,& cEnd); CHKERRQ(ierr);
  nelem = cEnd - cStart;

  // Get indices
  ierr = PetscMalloc1(nelem*P*P, &erestrict); CHKERRQ(ierr);
  for (c=cStart, eoffset = 0; c<cEnd; c++) {
    PetscInt numindices, *indices, i;
    ierr = DMPlexGetClosureIndices(dm, section, section, c, PETSC_TRUE,
                                   &numindices, &indices, NULL, NULL);
    CHKERRQ(ierr);
    for (i=0; i<numindices; i+=ncomp) {
      for (PetscInt j=0; j<ncomp; j++) {
        if (indices[i+j] != indices[i] + copysign(j, indices[i]))
          SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP,
                   "Cell %D closure indices not interlaced", c);
      }
      // NO BC on closed surfaces
      PetscInt loc = indices[i];
      erestrict[eoffset++] = loc;
    }
    ierr = DMPlexRestoreClosureIndices(dm, section, section, c, PETSC_TRUE,
                                       &numindices, &indices, NULL, NULL);
    CHKERRQ(ierr);
  }

  // Setup CEED restriction
  ierr = DMGetLocalVector(dm, &Uloc); CHKERRQ(ierr);
  ierr = VecGetLocalSize(Uloc, &nnodes); CHKERRQ(ierr);

  ierr = DMRestoreLocalVector(dm, &Uloc); CHKERRQ(ierr);
  CeedElemRestrictionCreate(ceed, nelem, P*P, ncomp, 1, nnodes, CEED_MEM_HOST,
                            CEED_COPY_VALUES, erestrict, Erestrict);
  ierr = PetscFree(erestrict); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// Set up libCEED for a given degree
static int SetupLibceedByDegree(DM dm, Ceed ceed, CeedInt degree,
                                CeedInt topodim, CeedInt qextra,
                                PetscInt ncompx, PetscInt ncompu,
                                PetscInt gsize, PetscInt xlsize,
                                bpType bpChoice, CeedData data,
                                PetscBool setup_rhs, CeedVector rhsceed,
                                CeedVector *target) {
  int ierr;
  DM dmcoord;
  Vec coords;
  const PetscScalar *coordArray;
  CeedBasis basisx, basisu;
  CeedElemRestriction Erestrictx, Erestrictu, Erestrictui, Erestrictqdi;
  CeedQFunction qf_setupgeo, qf_apply;
  CeedOperator op_setupgeo, op_apply;
  CeedVector xcoord, qdata, xceed, yceed;
  CeedInt P, Q, cStart, cEnd, nelem, qdatasize = bpOptions[bpChoice].qdatasize;
  CeedScalar R = 1,                      // radius of the sphere
             l = 1.0/PetscSqrtReal(3.0); // half edge of the inscribed cube

  // CEED bases
  P = degree + 1;
  Q = P + qextra;
  CeedBasisCreateTensorH1Lagrange(ceed, topodim, ncompu, P, Q,
                                  bpOptions[bpChoice].qmode, &basisu);
  CeedBasisCreateTensorH1Lagrange(ceed, topodim, ncompx, 2, Q,
                                  bpOptions[bpChoice].qmode, &basisx);

  // CEED restrictions
  ierr = DMGetCoordinateDM(dm, &dmcoord); CHKERRQ(ierr);
  ierr = DMPlexSetClosurePermutationTensor(dmcoord, PETSC_DETERMINE, NULL);
  CHKERRQ(ierr);

  ierr = CreateRestrictionPlex(ceed, 2, ncompx, &Erestrictx, dmcoord);
  CHKERRQ(ierr);
  ierr = CreateRestrictionPlex(ceed, P, ncompu, &Erestrictu, dm); CHKERRQ(ierr);

  ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd); CHKERRQ(ierr);
  nelem = cEnd - cStart;

  CeedElemRestrictionCreateStrided(ceed, nelem, Q*Q, ncompu, ncompu*nelem*Q*Q,
                                   CEED_STRIDES_BACKEND, &Erestrictui);
  CeedElemRestrictionCreateStrided(ceed, nelem, Q*Q, qdatasize,
                                   qdatasize*nelem*Q*Q,
                                   CEED_STRIDES_BACKEND, &Erestrictqdi);

  // Element coordinates
  ierr = DMGetCoordinatesLocal(dm, &coords); CHKERRQ(ierr);
  ierr = VecGetArrayRead(coords, &coordArray); CHKERRQ(ierr);

  CeedElemRestrictionCreateVector(Erestrictx, &xcoord, NULL);
  CeedVectorSetArray(xcoord, CEED_MEM_HOST, CEED_COPY_VALUES,
                     (PetscScalar *)coordArray);
  ierr = VecRestoreArrayRead(coords, &coordArray);

  // Create the persistent vectors that will be needed in setup and apply
  CeedInt nqpts;
  CeedBasisGetNumQuadraturePoints(basisu, &nqpts);
  CeedVectorCreate(ceed, qdatasize*nelem*nqpts, &qdata);
  CeedVectorCreate(ceed, xlsize, &xceed);
  CeedVectorCreate(ceed, xlsize, &yceed);

  // Create the Q-function that builds the operator (i.e. computes its
  // quadrature data) and set its context data
  CeedQFunctionCreateInterior(ceed, 1, bpOptions[bpChoice].setupgeo,
                              bpOptions[bpChoice].setupgeofname, &qf_setupgeo);
  CeedQFunctionAddInput(qf_setupgeo, "x", ncompx, CEED_EVAL_INTERP);
  CeedQFunctionAddInput(qf_setupgeo, "dx", ncompx*topodim, CEED_EVAL_GRAD);
  CeedQFunctionAddInput(qf_setupgeo, "weight", 1, CEED_EVAL_WEIGHT);
  CeedQFunctionAddOutput(qf_setupgeo, "qdata", qdatasize, CEED_EVAL_NONE);

  // Set up PDE operator
  CeedInt inscale = bpOptions[bpChoice].inmode==CEED_EVAL_GRAD ? topodim : 1;
  CeedInt outscale = bpOptions[bpChoice].outmode==CEED_EVAL_GRAD ? topodim : 1;
  CeedQFunctionCreateInterior(ceed, 1, bpOptions[bpChoice].apply,
                              bpOptions[bpChoice].applyfname, &qf_apply);
  CeedQFunctionAddInput(qf_apply, "u", ncompu*inscale,
                        bpOptions[bpChoice].inmode);
  CeedQFunctionAddInput(qf_apply, "qdata", qdatasize, CEED_EVAL_NONE);
  CeedQFunctionAddOutput(qf_apply, "v", ncompu*outscale,
                         bpOptions[bpChoice].outmode);

  // Create the operator that builds the quadrature data for the operator
  CeedOperatorCreate(ceed, qf_setupgeo, NULL, NULL, &op_setupgeo);
  CeedOperatorSetField(op_setupgeo, "x", Erestrictx, basisx,
                       CEED_VECTOR_ACTIVE);
  CeedOperatorSetField(op_setupgeo, "dx", Erestrictx, basisx,
                       CEED_VECTOR_ACTIVE);
  CeedOperatorSetField(op_setupgeo, "weight", CEED_ELEMRESTRICTION_NONE, basisx,
                       CEED_VECTOR_NONE);
  CeedOperatorSetField(op_setupgeo, "qdata", Erestrictqdi,
                       CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);

  // Create the mass or diff operator
  CeedOperatorCreate(ceed, qf_apply, NULL, NULL, &op_apply);
  CeedOperatorSetField(op_apply, "u", Erestrictu, basisu, CEED_VECTOR_ACTIVE);
  CeedOperatorSetField(op_apply, "qdata", Erestrictqdi, CEED_BASIS_COLLOCATED,
                       qdata);
  CeedOperatorSetField(op_apply, "v", Erestrictu, basisu, CEED_VECTOR_ACTIVE);

  // Setup qdata
  CeedOperatorApply(op_setupgeo, xcoord, qdata, CEED_REQUEST_IMMEDIATE);

  // Set up RHS if needed
  if (setup_rhs) {
    CeedQFunction qf_setuprhs;
    CeedOperator op_setuprhs;
    CeedVectorCreate(ceed, nelem*nqpts*ncompu, target);

    // Create the q-function that sets up the RHS and true solution
    CeedQFunctionCreateInterior(ceed, 1, bpOptions[bpChoice].setuprhs,
                                bpOptions[bpChoice].setuprhsfname, &qf_setuprhs);
    CeedQFunctionAddInput(qf_setuprhs, "x", ncompx, CEED_EVAL_INTERP);
    CeedQFunctionAddInput(qf_setuprhs, "qdata", qdatasize, CEED_EVAL_NONE);
    CeedQFunctionAddOutput(qf_setuprhs, "true_soln", ncompu, CEED_EVAL_NONE);
    CeedQFunctionAddOutput(qf_setuprhs, "rhs", ncompu, CEED_EVAL_INTERP);

    // Create the operator that builds the RHS and true solution
    CeedOperatorCreate(ceed, qf_setuprhs, NULL, NULL, &op_setuprhs);
    CeedOperatorSetField(op_setuprhs, "x", Erestrictx, basisx, CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_setuprhs, "qdata", Erestrictqdi, CEED_BASIS_COLLOCATED,
                         qdata);
    CeedOperatorSetField(op_setuprhs, "true_soln", Erestrictui,
                         CEED_BASIS_COLLOCATED,
                         *target);
    CeedOperatorSetField(op_setuprhs, "rhs", Erestrictu, basisu,
                         CEED_VECTOR_ACTIVE);

    // Set up the libCEED context
    CeedQFunctionContext rhsSetup;
    CeedQFunctionContextCreate(ceed, &rhsSetup);
    CeedScalar rhsSetupData[2] = {R, l};
    CeedQFunctionContextSetData(rhsSetup, CEED_MEM_HOST, CEED_COPY_VALUES,
                                sizeof rhsSetupData, &rhsSetupData);
    CeedQFunctionSetContext(qf_setuprhs, rhsSetup);
    CeedQFunctionContextDestroy(&rhsSetup);

    // Setup RHS and target
    CeedOperatorApply(op_setuprhs, xcoord, rhsceed, CEED_REQUEST_IMMEDIATE);
    CeedVectorTakeArray(rhsceed, CEED_MEM_HOST, NULL);

    // Cleanup
    CeedQFunctionDestroy(&qf_setuprhs);
    CeedOperatorDestroy(&op_setuprhs);
  }

  // Cleanup
  CeedQFunctionDestroy(&qf_setupgeo);
  CeedOperatorDestroy(&op_setupgeo);
  CeedVectorDestroy(&xcoord);

  // Save libCEED data required for level
  data->basisx = basisx; data->basisu = basisu;
  data->Erestrictx = Erestrictx;
  data->Erestrictu = Erestrictu;
  data->Erestrictui = Erestrictui;
  data->Erestrictqdi = Erestrictqdi;
  data->qf_apply = qf_apply;
  data->op_apply = op_apply;
  data->qdata = qdata;
  data->xceed = xceed;
  data->yceed = yceed;

  PetscFunctionReturn(0);
}

#ifdef multigrid
// Setup libCEED level transfer operator objects
static PetscErrorCode CeedLevelTransferSetup(Ceed ceed, CeedInt numlevels,
    CeedInt ncompu, bpType bpChoice, CeedData *data, CeedInt *leveldegrees,
    CeedQFunction qf_restrict, CeedQFunction qf_prolong) {
  // Return early if numlevels=1
  if (numlevels==1)
    PetscFunctionReturn(0);

  // Set up each level
  for (CeedInt i=1; i<numlevels; i++) {
    // P coarse and P fine
    CeedInt Pc = leveldegrees[i-1] + 1;
    CeedInt Pf = leveldegrees[i] + 1;

    // Restriction - Fine to corse
    CeedBasis basisctof;
    CeedOperator op_restrict;

    // Basis
    CeedBasisCreateTensorH1Lagrange(ceed, 3, ncompu, Pc, Pf,
                                    CEED_GAUSS_LOBATTO, &basisctof);

    // Create the restriction operator
    CeedOperatorCreate(ceed, qf_restrict, NULL, NULL, &op_restrict);
    CeedOperatorSetField(op_restrict, "input", data[i]->Erestrictu,
                         CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_restrict, "output", data[i-1]->Erestrictu,
                         basisctof, CEED_VECTOR_ACTIVE);

    // Save libCEED data required for level
    data[i]->basisctof = basisctof;
    data[i]->op_restrict = op_restrict;

    // Interpolation - Corse to fine
    CeedOperator op_interp;

    // Create the prolongation operator
    CeedOperatorCreate(ceed, qf_prolong, NULL, NULL, &op_interp);
    CeedOperatorSetField(op_interp, "uin", data[i-1]->Erestrictu,
                         basisctof, CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_interp, "uout", data[i]->Erestrictu,
                         CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);

    // Save libCEED data required for level
    data[i]->op_interp = op_interp;
  }

  PetscFunctionReturn(0);
}
#endif

// -----------------------------------------------------------------------------
// Mat Shell Functions
// -----------------------------------------------------------------------------

#ifdef multigrid
// This function returns the computed diagonal of the operator
static PetscErrorCode MatGetDiag(Mat A, Vec D) {
  PetscErrorCode ierr;
  UserO user;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  ierr = VecCopy(user->diag, D); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}
#endif

// This function uses libCEED to compute the action of the Laplacian with
// Dirichlet boundary conditions
static PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) {
  PetscErrorCode ierr;
  UserO user;
  PetscScalar *x, *y;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  // Global-to-local
  ierr = DMGlobalToLocalBegin(user->dm, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->dm, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = VecZeroEntries(user->Yloc); CHKERRQ(ierr);

  // Setup CEED vectors
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  ierr = VecGetArray(user->Yloc, &y); CHKERRQ(ierr);
  CeedVectorSetArray(user->xceed, CEED_MEM_HOST, CEED_USE_POINTER, x);
  CeedVectorSetArray(user->yceed, CEED_MEM_HOST, CEED_USE_POINTER, y);

  // Apply CEED operator
  CeedOperatorApply(user->op, user->xceed, user->yceed, CEED_REQUEST_IMMEDIATE);

  // Restore PETSc vectors
  CeedVectorTakeArray(user->xceed, CEED_MEM_HOST, NULL);
  CeedVectorTakeArray(user->yceed, CEED_MEM_HOST, NULL);
  ierr = VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x);
  CHKERRQ(ierr);
  ierr = VecRestoreArray(user->Yloc, &y); CHKERRQ(ierr);

  // Local-to-global
  ierr = VecZeroEntries(Y); CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(user->dm, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(user->dm, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

#ifdef multigrid
// This function uses libCEED to compute the action of the interp operator
static PetscErrorCode MatMult_Interp(Mat A, Vec X, Vec Y) {
  PetscErrorCode ierr;
  UserIR user;
  PetscScalar *x, *y;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  // Global-to-local
  ierr = VecZeroEntries(user->Xloc); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(user->dmc, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->dmc, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = VecZeroEntries(user->Yloc); CHKERRQ(ierr);

  // Setup CEED vectors
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  ierr = VecGetArray(user->Yloc, &y); CHKERRQ(ierr);
  CeedVectorSetArray(user->ceedvecc, CEED_MEM_HOST, CEED_USE_POINTER, x);
  CeedVectorSetArray(user->ceedvecf, CEED_MEM_HOST, CEED_USE_POINTER, y);

  // Apply CEED operator
  CeedOperatorApply(user->op, user->ceedvecc, user->ceedvecf,
                    CEED_REQUEST_IMMEDIATE);

  // Restore PETSc vectors
  CeedVectorTakeArray(user->ceedvecc, CEED_MEM_HOST, NULL);
  CeedVectorTakeArray(user->ceedvecf, CEED_MEM_HOST, NULL);
  ierr = VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x);
  CHKERRQ(ierr);
  ierr = VecRestoreArray(user->Yloc, &y); CHKERRQ(ierr);

  // Multiplicity
  ierr = VecPointwiseMult(user->Yloc, user->Yloc, user->mult);

  // Local-to-global
  ierr = VecZeroEntries(Y); CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(user->dmf, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(user->dmf, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// This function uses libCEED to compute the action of the restriction operator
static PetscErrorCode MatMult_Restrict(Mat A, Vec X, Vec Y) {
  PetscErrorCode ierr;
  UserIR user;
  PetscScalar *x, *y;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  // Global-to-local
  ierr = VecZeroEntries(user->Xloc); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(user->dmf, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->dmf, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = VecZeroEntries(user->Yloc); CHKERRQ(ierr);

  // Multiplicity
  ierr = VecPointwiseMult(user->Xloc, user->Xloc, user->mult); CHKERRQ(ierr);

  // Setup CEED vectors
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  ierr = VecGetArray(user->Yloc, &y); CHKERRQ(ierr);
  CeedVectorSetArray(user->ceedvecf, CEED_MEM_HOST, CEED_USE_POINTER, x);
  CeedVectorSetArray(user->ceedvecc, CEED_MEM_HOST, CEED_USE_POINTER, y);

  // Apply CEED operator
  CeedOperatorApply(user->op, user->ceedvecf, user->ceedvecc,
                    CEED_REQUEST_IMMEDIATE);

  // Restore PETSc vectors
  CeedVectorTakeArray(user->ceedvecf, CEED_MEM_HOST, NULL);
  CeedVectorTakeArray(user->ceedvecc, CEED_MEM_HOST, NULL);
  ierr = VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x);
  CHKERRQ(ierr);
  ierr = VecRestoreArray(user->Yloc, &y); CHKERRQ(ierr);

  // Local-to-global
  ierr = VecZeroEntries(Y); CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(user->dmc, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(user->dmc, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);

  PetscFunctionReturn(0);
}
#endif

// This function calculates the error in the final solution
static PetscErrorCode ComputeErrorMax(UserO user, CeedOperator op_error,
                                      Vec X, CeedVector target,
                                      PetscReal *maxerror) {
  PetscErrorCode ierr;
  PetscScalar *x;
  CeedVector collocated_error;
  CeedInt length;

  PetscFunctionBeginUser;
  CeedVectorGetLength(target, &length);
  CeedVectorCreate(user->ceed, length, &collocated_error);

  // Global-to-local
  ierr = DMGlobalToLocal(user->dm, X, INSERT_VALUES, user->Xloc); CHKERRQ(ierr);

  // Setup CEED vector
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  CeedVectorSetArray(user->xceed, CEED_MEM_HOST, CEED_USE_POINTER, x);

  // Apply CEED operator
  CeedOperatorApply(op_error, user->xceed, collocated_error,
                    CEED_REQUEST_IMMEDIATE);

  // Restore PETSc vector
  VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);

  // Reduce max error
  *maxerror = 0;
  const CeedScalar *e;
  CeedVectorGetArrayRead(collocated_error, CEED_MEM_HOST, &e);
  for (CeedInt i=0; i<length; i++) {
    *maxerror = PetscMax(*maxerror, PetscAbsScalar(e[i]));
  }
  CeedVectorRestoreArrayRead(collocated_error, &e);
  ierr = MPI_Allreduce(MPI_IN_PLACE, maxerror,
                       1, MPIU_REAL, MPIU_MAX, user->comm); CHKERRQ(ierr);

  // Cleanup
  CeedVectorDestroy(&collocated_error);

  PetscFunctionReturn(0);
}

#endif // setupsphere_h