#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function

"""
Empirical XMM background model.
Written by Zhu Liu, adapted by Torben Simm and Johannes Buchner
(C) 2013-2016
For example usage, see examples/sherpa/background/xmm/fit.py
"""

import os
import logging

import numpy

if "MAKESPHINXDOC" not in os.environ:
    import sherpa.astro.ui as ui
    from sherpa.stats import Cash, CStat
    from sherpa.models.parameter import Parameter
    from sherpa.models import ArithmeticModel, CompositeModel
    from sherpa.astro.ui import *
    from sherpa.astro.instrument import RSPModelNoPHA, RMFModelNoPHA
else:
    CompositeModel, ArithmeticModel = object, object

from bxa.sherpa.background.xmm import get_embedded_file

# print(
#     """
# Using XMM empirical background model originally by Richard Sturm.
# Please reference Maggi P., et al., 2014, A&A, 561, AA76.
# """
# )

# def get_embedded_file(filename):
#     """
#     Gets the path of a file in the same folder as this script
#     """
#     return os.path.join(os.path.dirname(__file__), filename)

def get_unitrmf(rmf):
    """
    Returns a unit response matrix with the same properties
    (channels, energies, etc) as the original rmf
    """
    urmf = unpack_rmf(rmf.name)
    urmf.matrix = numpy.array([1.0] * urmf.detchans)
    urmf.f_chan = numpy.array(range(urmf.detchans), dtype=numpy.uint32)
    urmf.n_chan = numpy.array([1] * urmf.detchans, dtype=numpy.uint32)
    urmf.offset = 0
    urmf._rsp = numpy.array([1.0] * urmf.detchans)
    urmf._fch = numpy.array(range(urmf.detchans), dtype=numpy.uint32)
    urmf._nch = numpy.array([1] * urmf.detchans, dtype=numpy.uint32)

    return urmf

def get_pn_bkg_model(i, galabs, fit=False, fix_response=False):
    # =================================================================
    # Parameters:
    # i = sherpa ID of data set
    # galabs = name of model for galactic absorption
    # pnbrsp = response model of bkg data set
    # fit: True = fit bkg spectrum; False = just set_bkg_full_model
    # returns unit response for bkg model
    # =================================================================
    # get instrument response model for bkg spectrum
    if fix_response:
        if get_rmf(i).energ_lo[0] == 0:
            get_rmf(i).energ_lo[0] = 0.001

        if get_arf(i).energ_lo[0] == 0:
            get_arf(i).energ_lo[0] = 0.001

    pnbrsp = get_response(i, bkg_id=1)
    pnscale = get_bkg_scale(i)  # get background scaling factor

    # create unit response
    dia_pn_rmf=get_embedded_file('pn_dia.rmf')
    dia_pn_arf=get_embedded_file('pn_dia.arf')
    copy_data(i,1000+2)
    load_bkg_rmf(1000+2, dia_pn_rmf) #load diagonal bkg matrices
    load_bkg_arf(1000+2, dia_pn_arf)
    pnbunitrsp = get_response(1000+2, bkg_id=1)
    delete_data(1000+2)

    # gaussian line center energy, line width, and initial normalization; for *PN background*
    pncenters = [
        1.49165,
        1.49165,
        4.53177,
        5.42516,
        6.38155,
        7.48675,
        8.04087,
        8.04087,
        8.60924,
        8.89395,
        9.56160,
    ]
    pnlinewidth = [
        5.73813e-02,
        3.63469e-02,
        6.10487e-02,
        7.08380e-02,
        9.59053e-02,
        6.52422e-02,
        9.48594e-02,
        6.26174e-05,
        0.120893,
        0.114254,
        0.108717,
    ]
    pnlinenorm = [
        7.81356e-03,
        3.96601e-03,
        7.30727e-04,
        4.96413e-04,
        5.31295e-9,
        6.84796e-04,
        3.01564e-02,
        1.41847e-04,
        8.87887e-03,
        5.75592e-03,
        1.71367e-03,
    ]

    # model component prefix
    bkg_prefix = "bkg{}_pn".format(i)

    (
        pnbkgcons,
        pnbkgspline1,
        pnbkgexpdec,
        pnbkgsmedge1,
        pnbkgsmedge2,
        pnbkgspline2,
        pnbkginspl,
        pnbkgline1,
        pnbkgline2,
        pnbkgline3,
        pnbkgline4,
        pnbkgline5,
        pnbkgline6,
        pnbkgline7,
        pnbkgline8,
        pnbkgline9,
        pnbkgline10,
        pnbkgline11,
        pnbkgpl,
        pnbkgapec,
        pnbkglcapec,
    ) = (
        xsconstant(bkg_prefix + "cons"),
        xsspline(bkg_prefix + "spline1"),
        xsexpdec(bkg_prefix + "expdec"),
        xssmedge(bkg_prefix + "smedge1"),
        xssmedge(bkg_prefix + "smedge2"),
        xsspline(bkg_prefix + "spline2"),
        xspowerlaw(bkg_prefix + "bkpl"),
        xsgaussian(bkg_prefix + "gau1"),
        xsgaussian(bkg_prefix + "gau2"),
        xsgaussian(bkg_prefix + "gau3"),
        xsgaussian(bkg_prefix + "gau4"),
        xsgaussian(bkg_prefix + "gau5"),
        xsgaussian(bkg_prefix + "gau6"),
        xsgaussian(bkg_prefix + "gau7"),
        xsgaussian(bkg_prefix + "gau8"),
        xsgaussian(bkg_prefix + "gau9"),
        xsgaussian(bkg_prefix + "gau10"),
        xsgaussian(bkg_prefix + "gau11"),
        xspowerlaw(bkg_prefix + "expl"),
        xsapec(bkg_prefix + "apec"),
        xsapec(bkg_prefix + "lcapec"),
    )

    pnlines = [
        pnbkgline1,
        pnbkgline2,
        pnbkgline3,
        pnbkgline4,
        pnbkgline5,
        pnbkgline6,
        pnbkgline7,
        pnbkgline8,
        pnbkgline9,
        pnbkgline10,
        pnbkgline11,
    ]
    pnfixwid = [
        pnbkgline2,
        pnbkgline3,
        pnbkgline4,
        pnbkgline5,
        pnbkgline6,
        pnbkgline8,
        pnbkgline11,
    ]
    pnfree = [pnbkgline1, pnbkgline7, pnbkgline9, pnbkgline10]

    # define PN background model
    pn_bkg = pnbunitrsp(
        pnbkgcons
        * (
            pnbkgspline1 * pnbkgexpdec
            + pnbkgsmedge1
            * pnbkgsmedge2
            * (
                pnbkgspline2 * pnbkginspl
                + pnbkgline1
                + pnbkgline2
                + pnbkgline3
                + pnbkgline4
                + pnbkgline5
                + pnbkgline6
                + pnbkgline7
                + pnbkgline8
                + pnbkgline9
                + pnbkgline10
                + pnbkgline11
            )
        )
    ) + pnbrsp(galabs * (pnbkgpl + pnbkgapec) + pnbkglcapec)

    for l, c in zip(pnlines, pncenters):
        l.LineE = c
        l.LineE.min = c - 0.05
        l.LineE.max = c + 0.05

    pnbkgline2.LineE = pnbkgline1.LineE
    pnbkgline8.LineE = pnbkgline7.LineE

    for l, s in zip(pnlines, pnlinewidth):
        l.Sigma = s

    for l in pnfree:
        l.Sigma.min = 1e-5
        l.Sigma.max = 0.2

    for l in pnfixwid:
        l.Sigma.freeze()

    for l, n in zip(pnlines, pnlinenorm):
        l.norm = n
        l.norm.min = 1e-10
        l.norm.max = 1e10

    # Scaling constant
    pnbkgcons.factor = 1.0
    pnbkgcons.factor.freeze()

    # thermal radiation
    pnbkgapec.kT = 0.286928
    pnbkgapec.kT.min = 0.008
    pnbkgapec.kT.max = 64
    pnbkgapec.Abundanc = 1.0
    pnbkgapec.Abundanc.freeze()
    pnbkgapec.Redshift = 0.0
    pnbkgapec.Redshift.freeze()
    pnbkgapec.norm = 5.58410e-05
    pnbkgapec.norm.min = 1e-10
    pnbkgapec.norm.max = 1e10

    # local thermal radiation
    pnbkglcapec.kT = 0.1
    pnbkglcapec.kT.freeze()
    pnbkglcapec.Abundanc = 1.0
    pnbkglcapec.Abundanc.freeze()
    pnbkglcapec.Redshift = 0.0
    pnbkglcapec.Redshift.freeze()
    pnbkglcapec.norm = 3.89164e-05
    pnbkglcapec.norm.min = 1e-10
    pnbkglcapec.norm.max = 1e10

    # exponential decay
    pnbkgexpdec.factor = 44.3418
    pnbkgexpdec.factor.min = 0
    pnbkgexpdec.factor.max = 100
    pnbkgexpdec.norm = 6830.89
    pnbkgexpdec.norm.freeze()

    # Smear function
    pnbkgsmedge1.edgeE = 0.538408
    pnbkgsmedge1.edgeE.freeze()
    pnbkgsmedge1.MaxTau = 1.40238
    pnbkgsmedge1.MaxTau.min = 0
    pnbkgsmedge1.MaxTau.max = 10
    pnbkgsmedge1.index = -2.67000
    pnbkgsmedge1.index.freeze()
    pnbkgsmedge1.width = 0.313365
    pnbkgsmedge1.width.min = 0.01
    pnbkgsmedge1.width.max = 100

    pnbkgsmedge2.edgeE = 1.38826
    pnbkgsmedge2.edgeE.freeze()
    pnbkgsmedge2.MaxTau.min = 0
    pnbkgsmedge2.MaxTau.max = 10
    pnbkgsmedge2.MaxTau = 9.37167
    pnbkgsmedge2.index = -2.67000
    pnbkgsmedge2.index.freeze()
    pnbkgsmedge2.width = 5.7642
    pnbkgsmedge2.width.min = 0.01
    pnbkgsmedge2.width.max = 100

    # Spline funtion
    pnbkgspline1.Estart = 0.200000
    pnbkgspline1.Estart.freeze()
    pnbkgspline1.Ystart = -1.31506
    pnbkgspline1.Ystart.min = -1e6
    pnbkgspline1.Ystart.max = 1e6
    pnbkgspline1.Yend = 1064.16
    pnbkgspline1.Yend.min = -1e6
    pnbkgspline1.Yend.max = 1e6
    pnbkgspline1.YPstart = -106.183
    pnbkgspline1.YPstart.min = -1e6
    pnbkgspline1.YPstart.max = 1e6
    pnbkgspline1.YPend = -366.092
    pnbkgspline1.YPend.min = -1e6
    pnbkgspline1.YPend.max = 1e6
    pnbkgspline1.Eend = 1.74715
    pnbkgspline1.Eend.min = 0
    pnbkgspline1.Eend.max = 100

    pnbkgspline2.Estart = 3.29056
    pnbkgspline2.Ystart = 1.00643
    pnbkgspline2.Ystart.min = -1e6
    pnbkgspline2.Ystart.max = 1e6
    pnbkgspline2.Yend = 0.887026
    pnbkgspline2.Yend.min = -1e6
    pnbkgspline2.Yend.max = 1e6
    pnbkgspline2.YPstart = -0.278401
    pnbkgspline2.YPstart.min = -1e6
    pnbkgspline2.YPstart.max = 1e6
    pnbkgspline2.YPend = 4.84809e-03
    pnbkgspline2.YPend.min = -1e6
    pnbkgspline2.YPend.max = 1e6
    pnbkgspline2.Eend = 7.32701
    pnbkgspline2.Eend.min = 0
    pnbkgspline2.Eend.max = 100

    # Background power law
    pnbkginspl.PhoIndex = 0.279
    pnbkginspl.PhoIndex.min = -2
    pnbkginspl.PhoIndex.max = 9
    pnbkginspl.norm = 8.23614e-03
    pnbkginspl.norm.min = 1e-10
    pnbkginspl.norm.max = 1e6

    # Extragalactic background
    pnbkgpl.PhoIndex = 1.46
    pnbkgpl.PhoIndex.freeze()
    pnbkgpl.norm = 1.25288e-05
    pnbkgpl.norm.min = 1e-10
    pnbkgpl.norm.max = 1e3

    if fit:
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1 * pnbkgsmedge2 * (pnbkgspline2 * pnbkginspl)
                )
            ),
        )
        logging.info("Fitting (1/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (pnbkgspline2 * pnbkginspl + pnbkgline2)
                )
            ),
        )
        logging.info("Fitting (2/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (pnbkgspline2 * pnbkginspl + pnbkgline1 + pnbkgline2)
                )
            ),
        )
        logging.info("Fitting (3/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline7
                        + pnbkgline8
                    )
                )
            ),
        )
        logging.info("Fitting (4/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline7
                        + pnbkgline8
                        + pnbkgline9
                        + pnbkgline10
                    )
                )
            ),
        )
        logging.info("Fitting (5/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline3
                        + pnbkgline4
                        + pnbkgline5
                        + pnbkgline6
                        + pnbkgline7
                        + pnbkgline8
                        + pnbkgline9
                        + pnbkgline10
                        + pnbkgline11
                    )
                )
            ),
        )
        logging.info("Fitting (6/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline3
                        + pnbkgline4
                        + pnbkgline5
                        + pnbkgline6
                        + pnbkgline7
                        + pnbkgline8
                        + pnbkgline9
                        + pnbkgline10
                        + pnbkgline11
                    )
                )
            )
            + pnbrsp(galabs * (pnbkgapec)),
        )
        logging.info("Fitting (7/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline3
                        + pnbkgline4
                        + pnbkgline5
                        + pnbkgline6
                        + pnbkgline7
                        + pnbkgline8
                        + pnbkgline9
                        + pnbkgline10
                        + pnbkgline11
                    )
                )
            )
            + pnbrsp(galabs * (pnbkgapec + pnbkgpl)),
        )
        logging.info("Fitting (8/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline3
                        + pnbkgline4
                        + pnbkgline5
                        + pnbkgline6
                        + pnbkgline7
                        + pnbkgline8
                        + pnbkgline9
                        + pnbkgline10
                        + pnbkgline11
                    )
                )
            )
            + pnbrsp(galabs * (pnbkgapec + pnbkgpl) + pnbkglcapec),
        )
        fit_bkg(i)

        freeze(
            pnbkgcons,
            pnbkgspline1,
            pnbkgexpdec,
            pnbkgsmedge1,
            pnbkgsmedge2,
            pnbkgspline2,
            pnbkginspl,
            pnbkgline1,
            pnbkgline2,
            pnbkgline3,
            pnbkgline4,
            pnbkgline5,
            pnbkgline6,
            pnbkgline7,
            pnbkgline8,
            pnbkgline9,
            pnbkgline10,
            pnbkgline11,
            galabs,
            pnbkgapec,
            pnbkgpl,
            pnbkglcapec,
        )
        logging.info("PN background model set up and fitted")
        logging.info("Please double-check that it is a good fit")
    else:
        set_bkg_full_model(
            i,
            pnbunitrsp(
                pnbkgcons
                * (
                    pnbkgspline1 * pnbkgexpdec
                    + pnbkgsmedge1
                    * pnbkgsmedge2
                    * (
                        pnbkgspline2 * pnbkginspl
                        + pnbkgline1
                        + pnbkgline2
                        + pnbkgline3
                        + pnbkgline4
                        + pnbkgline5
                        + pnbkgline6
                        + pnbkgline7
                        + pnbkgline8
                        + pnbkgline9
                        + pnbkgline10
                        + pnbkgline11
                    )
                )
            )
            + pnbrsp(galabs * (pnbkgapec + pnbkgpl) + pnbkglcapec),
        )
        logging.info("PN background model set up")

    return pnscale * (
        pnbunitrsp(
            pnbkgcons
            * (
                pnbkgspline1 * pnbkgexpdec
                + pnbkgsmedge1
                * pnbkgsmedge2
                * (
                    pnbkgspline2 * pnbkginspl
                    + pnbkgline1
                    + pnbkgline2
                    + pnbkgline3
                    + pnbkgline4
                    + pnbkgline5
                    + pnbkgline6
                    + pnbkgline7
                    + pnbkgline8
                    + pnbkgline9
                    + pnbkgline10
                    + pnbkgline11
                )
            )
        )
        + pnbrsp(galabs * (pnbkgapec + pnbkgpl) + pnbkglcapec)
    )

def get_mos_bkg_model(i, galabs, fit=False, fix_response=False):
    # =================================================================
    # Parameters:
    # i = sherpa ID of data set
    # galabs = name of model for galactic absorption
    # fit: True = fit bkg spectrum; False = just set_bkg_full_model
    # returns unit response for bkg model
    # =================================================================
    # get instrument response model for bkg spectrum
    if fix_response:
        if get_rmf(i).energ_lo[0] == 0:
            get_rmf(i).energ_lo[0] = 0.001
        if get_arf(i).energ_lo[0] == 0:
            get_arf(i).energ_lo[0] = 0.001

    mosbrsp = get_response(i, bkg_id=1)
    mosscale = get_bkg_scale(i)  # get background scaling factor

    # create unit response
    dia_mos_rmf=get_embedded_file('mos_dia.rmf')
    dia_mos_arf=get_embedded_file('mos_dia.arf')
    copy_data(i,1000+2)
    load_bkg_rmf(1000+2, dia_mos_rmf) #load diagonal bkg matrices
    load_bkg_arf(1000+2, dia_mos_arf)
    mosbunitrsp = get_response(1000+2, bkg_id=1)
    delete_data(1000+2)

    # gaussian line center energy, line width, and initial normalization; for *MOS background*
    moscenters = [1.48600, 1.48700, 1.74000, 5.41000, 5.89500, 6.42000, 9.71000]
    moslinewidth = [
        3.84602e-02,
        0.165816,
        3.54985e-02,
        9.77018e-02,
        7.45076e-02,
        7.42365e-02,
        9.04855e-02,
    ]
    moslinenorm = [
        9.93119e-03,
        1.67028e-03,
        1.75461e-03,
        2.86358e-04,
        2.07525e-04,
        3.07555e-04,
        4.58115e-04,
    ]

    # model component prefix
    bkg_prefix = "bkg{}_mos".format(i)

    (
        mosbkgcons,
        mosbkgsmedge,
        mosbkgspline,
        mosbkgbknpl,
        mosbkgline1,
        mosbkgline2,
        mosbkgline3,
        mosbkgline4,
        mosbkgline5,
        mosbkgline6,
        mosbkgline7,
        mosbkgpl,
        mosbkgapec,
        mosbkglcapec,
    ) = (
        xsconstant(bkg_prefix + "cons"),
        xssmedge(bkg_prefix + "smedge"),
        xsspline(bkg_prefix + "spline"),
        xsbknpower(bkg_prefix + "bknpl"),
        xsgaussian(bkg_prefix + "gau1"),
        xsgaussian(bkg_prefix + "gau2"),
        xsgaussian(bkg_prefix + "gau3"),
        xsgaussian(bkg_prefix + "gau4"),
        xsgaussian(bkg_prefix + "gau5"),
        xsgaussian(bkg_prefix + "gau6"),
        xsgaussian(bkg_prefix + "gau7"),
        xspowerlaw(bkg_prefix + "expl"),
        xsapec(bkg_prefix + "apec"),
        xsapec(bkg_prefix + "lcapec"),
    )

    moslines = [
        mosbkgline1,
        mosbkgline2,
        mosbkgline3,
        mosbkgline4,
        mosbkgline5,
        mosbkgline6,
        mosbkgline7,
    ]

    # define MOS background model
    mos_bkg = mosbunitrsp(
        mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
        + mosbkgline1
        + mosbkgline2
        + mosbkgline3
        + mosbkgline4
        + mosbkgline5
        + mosbkgline6
        + mosbkgline7
    ) + mosbrsp(galabs * (mosbkgpl + mosbkgapec) + mosbkglcapec)

    for l, c in zip(moslines, moscenters):
        l.LineE = c
        l.LineE.min = c - 0.05
        l.LineE.max = c + 0.05

    for l, s in zip(moslines, moslinewidth):
        l.Sigma = s
        l.Sigma.min = 1e-4
        l.Sigma.max = 0.2
    mosbkgline1.Sigma.min = 1e-4
    mosbkgline1.Sigma.max = 1e-1

    for l, n in zip(moslines, moslinenorm):
        l.norm = n
        l.norm.min = 1e-10
        l.norm.max = 1e10

    # Constant factor
    mosbkgcons.factor = 1.0
    mosbkgcons.factor.freeze()

    # Smear function
    mosbkgsmedge.edgeE = 0.538408
    mosbkgsmedge.edgeE.freeze()
    mosbkgsmedge.MaxTau = 0.246633
    mosbkgsmedge.MaxTau.min = 0.0
    mosbkgsmedge.MaxTau.max = 10.0
    mosbkgsmedge.index = -2.67
    mosbkgsmedge.index.freeze()
    mosbkgsmedge.width = 1e-02
    mosbkgsmedge.width.min = 1e-02
    mosbkgsmedge.width.max = 1e2

    # Spline funtion
    mosbkgspline.Estart = 3.08175
    mosbkgspline.Ystart = 1.00984
    mosbkgspline.Yend = 1.99144
    mosbkgspline.YPstart = -2.90195e-02
    mosbkgspline.YPend = 5.49102e-02
    mosbkgspline.Estart.freeze()
    mosbkgspline.Ystart.freeze()
    mosbkgspline.Yend.freeze()
    mosbkgspline.YPstart.freeze()
    mosbkgspline.YPend.freeze()
    mosbkgspline.Eend = 13.6492
    mosbkgspline.Eend.min = 0
    mosbkgspline.Eend.max = 100

    # Broken power law
    mosbkgbknpl.PhoIndx1 = 1.48636
    mosbkgbknpl.PhoIndx1.min = -2
    mosbkgbknpl.PhoIndx1.max = 9
    mosbkgbknpl.BreakE = 0.415173
    mosbkgbknpl.PhoIndx2 = 0.315615
    mosbkgbknpl.BreakE.freeze()
    mosbkgbknpl.PhoIndx2.freeze()
    mosbkgbknpl.norm = 2.90071e-03
    mosbkgbknpl.norm.min = 1e-10
    mosbkgbknpl.norm.max = 1e10

    # Extragalactic background
    mosbkgpl.PhoIndex = 1.46
    mosbkgpl.PhoIndex.freeze()
    mosbkgpl.norm = 4.58115e-04
    mosbkgpl.norm.min = 1e-10
    mosbkgpl.norm.max = 1e10

    # thermal radiation
    mosbkgapec.kT = 0.286928
    mosbkgapec.kT.min = 0.008
    mosbkgapec.kT.max = 64
    mosbkgapec.Abundanc = 1.0
    mosbkgapec.Abundanc.freeze()
    mosbkgapec.Redshift = 0.0
    mosbkgapec.Redshift.freeze()
    mosbkgapec.norm = 5.58410e-05
    mosbkgapec.norm.min = 1e-10
    mosbkgapec.norm.max = 1e10

    # local thermal radiation
    mosbkglcapec.kT = 0.1
    mosbkglcapec.kT.freeze()
    mosbkglcapec.Abundanc = 1.0
    mosbkglcapec.Abundanc.freeze()
    mosbkglcapec.Redshift = 0.0
    mosbkglcapec.Redshift.freeze()
    mosbkglcapec.norm = 3.89164e-05
    mosbkglcapec.norm.min = 1e-10
    mosbkglcapec.norm.max = 1e10

    if fit:  # fit bkg
        set_bkg_full_model(
            i, mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
            )
        )
        logging.info("Fitting (1/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline2
            ),
        )
        logging.info("Fitting (2/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline2
                + mosbkgline3
            ),
        )
        logging.info("Fitting (3/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline1
                + mosbkgline2
                + mosbkgline3
            ),
        )
        logging.info("Fitting (4/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline1
                + mosbkgline2
                + mosbkgline3
                + mosbkgline4
                + mosbkgline5
                + mosbkgline6
                + mosbkgline7
            ),
        )
        logging.info("Fitting (5/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline1
                + mosbkgline2
                + mosbkgline3
                + mosbkgline4
                + mosbkgline5
                + mosbkgline6
                + mosbkgline7
            )
            + mosbrsp(galabs * (mosbkgapec)),
        )
        logging.info("Fitting (6/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline1
                + mosbkgline2
                + mosbkgline3
                + mosbkgline4
                + mosbkgline5
                + mosbkgline6
                + mosbkgline7
            )
            + mosbrsp(galabs * (mosbkgapec + mosbkgpl)),
        )
        logging.info("Fitting (7/8)...")
        fit_bkg(i)
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline1
                + mosbkgline2
                + mosbkgline3
                + mosbkgline4
                + mosbkgline5
                + mosbkgline6
                + mosbkgline7
            )
            + mosbrsp(galabs * (mosbkgapec + mosbkgpl) + mosbkglcapec),
        )
        logging.info("Fitting (8/8)...")
        fit_bkg(i)
        freeze(
            mosbkgcons,
            mosbkgsmedge,
            mosbkgspline,
            mosbkgbknpl,
            mosbkgline1,
            mosbkgline2,
            mosbkgline3,
            mosbkgline4,
            mosbkgline5,
            mosbkgline6,
            mosbkgline7,
            mosbkgpl,
            mosbkgapec,
            mosbkglcapec,
        )
        logging.info("MOS background model set up and fitted")
        logging.info("Please double-check that it is a good fit")
    else:
        set_bkg_full_model(
            i,
            mosbunitrsp(
                mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
                + mosbkgline1
                + mosbkgline2
                + mosbkgline3
                + mosbkgline4
                + mosbkgline5
                + mosbkgline6
                + mosbkgline7
            )
            + mosbrsp(galabs * (mosbkgapec + mosbkgpl) + mosbkglcapec),
        )
        logging.info("MOS background model set up")

    return mosscale * (
        mosbunitrsp(
            mosbkgcons * mosbkgsmedge * mosbkgspline * mosbkgbknpl
            + mosbkgline1
            + mosbkgline2
            + mosbkgline3
            + mosbkgline4
            + mosbkgline5
            + mosbkgline6
            + mosbkgline7
        )
        + mosbrsp(galabs * (mosbkgapec + mosbkgpl) + mosbkglcapec)
    )

def get_mos_bkg_model_cached(i, galabs):
    filename = get_bkg(i).name + ".bkgpars"
    if os.path.exists(filename):
        bkgmodel = get_mos_bkg_model(i, galabs, fit=False)
        for p, v in zip(bkgmodel.pars, numpy.loadtxt(filename)):
            p.val = v
    else:
        bkgmodel = get_mos_bkg_model(i, galabs, fit=True)
        numpy.savetxt(filename, [p.val for p in bkgmodel.pars])
    for p in bkgmodel.pars:
        p.freeze()

    return bkgmodel

def get_pn_bkg_model_cached(i, galabs):
    filename = get_bkg(i).name + ".bkgpars"
    if os.path.exists(filename):
        bkgmodel = get_pn_bkg_model(i, galabs, fit=False)
        for p, v in zip(bkgmodel.pars, numpy.loadtxt(filename)):
            p.val = v
    else:
        bkgmodel = get_pn_bkg_model(i, galabs, fit=True)
        numpy.savetxt(filename, [p.val for p in bkgmodel.pars])
    for p in bkgmodel.pars:
        p.freeze()

    return bkgmodel


#!/usr/bin/env python

"""
specter.throughput
==================

A class for tracking throughput.

Hacks:

- Doesn't support  spatial variation of input sources.
- Doesn't support per-fiber throughput.
- Do I really want to impose a clumsy ObjType ENUM?

How to handle fiber size and sky units?
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import sys
import os
import warnings
import numbers
import numpy as np
from astropy.io import fits
from specter import util

#- ObjType enum
class ObjType:
    STAR   = 'STAR'
    STD    = 'STD'
    GALAXY = 'GALAXY'   #- pretty meaningless for fiberloss, ok for others
    ELG    = 'ELG'
    LRG    = 'LRG'
    QSO    = 'QSO'
    SKY    = 'SKY'
    CALIB  = 'CALIB'

def load_throughput(filename):
    """
    Create Throughput object from FITS file with EXTNAME=THROUGHPUT HDU
    """
    #- memmap=False so that fits will really close the file upon fx.close()
    fx = fits.open(filename, memmap=False)
    thru = fx['THROUGHPUT'].data
    hdr =  fx['THROUGHPUT'].header

    #- Check for FIBERINPUT HDU
    if 'FIBERINPUT' in fx:
        tmp = fx['FIBERINPUT'].data
        assert(len(tmp) == len(thru))
        fiberinput = dict()
        for key in tmp.dtype.names:
            fiberinput[key.upper()] = tmp[key]
    else:
        print("no FIBERINPUT extention found")
        fiberinput = thru['fiberinput']

    if 'wavelength' in thru.dtype.names:
        w = thru['wavelength']
    elif 'loglam' in thru.dtype.names:
        w = 10**thru['loglam']
    else:
        fx.close()
        raise ValueError('throughput must include wavelength or loglam')

    if 'GEOMAREA' in hdr:
        area = hdr['GEOMAREA']
    elif 'EFFAREA' in hdr:
        area = hdr['EFFAREA']  #- misnomer, but for backwards compatibility
    elif 'AREA' in hdr:
        area = hdr['AREA']
    else:
        fx.close()
        raise ValueError("throughput file missing GEOMAREA keyword")

    fx.close()

    return Throughput(
        wave = w,
        throughput = thru['throughput'],
        extinction = thru['extinction'],
        fiberinput = fiberinput,
        exptime    = hdr['EXPTIME'],
        area       = area,
        fiberdia   = hdr['FIBERDIA'],
        )

class Throughput:
    def __init__(self, wave, throughput, extinction,
        exptime, area, fiberdia, fiberinput=None):
        """
        Create Throughput object

        Inputs
        ------
        wave : wavelength array [Angstroms]
        throughput : array of system throughput for elements which apply to
            all types of sources, e.g. mirrors, lenses, CCD efficiency
        extinction : atmospheric extinction array [mags/airmass]

        exptime:  float, default exposure time [sec]
        area:     float, input geometric area [cm^2]
        fiberdia: float, fiber diameter [arcsec]

        Optional Inputs
        ---------------
        fiberinput : float, array, or dictionary of arrays keyed by objtype.
            Geometric throughput due to finite sized fiber input.
            Default to no loss = 1.0.

        Notes
        -----
        fiberinput is a placeholder, since it really depends upon the
        spatial extent of the object and the seeing.
        """
        self._wave = np.copy(wave)
        self._thru = np.copy(throughput)
        self._extinction  = np.copy(extinction)

        self.exptime = float(exptime)
        self.area = float(area)
        self.fiberdia = float(fiberdia)

        #- Flux -> photons conversion constant
        #-         h [erg s]      * c [m/s]      * [1e10 A/m] = [erg A]
        self._hc = 6.62606957e-27 * 2.99792458e8 * 1e10

        #- Create fiber input dict keyed by object type, including 'default'
        if fiberinput is not None:
            if isinstance(fiberinput, numbers.Real):
                self._fiberinput = dict(default=np.ones(len(wave)) * fiberinput)
            elif isinstance(fiberinput, np.ndarray):
                self._fiberinput = dict(default=np.copy(fiberinput))
            elif isinstance(fiberinput, dict):
                self._fiberinput = fiberinput
                if 'default' not in fiberinput:
                    self._fiberinput['default'] = np.ones(len(wave))
            else:
                raise ValueError('Unrecognized type for fiberinput: {}'.format(type(fiberinput)))
        else:
            self._fiberinput = dict(default=np.ones(len(wave)))

        #- special cases: QSO and STD are STAR for fiber input losses
        if 'STAR' in self._fiberinput and 'STD' not in self._fiberinput:
            self._fiberinput['STD'] = self._fiberinput['STAR']
        if 'STAR' in self._fiberinput and 'QSO' not in self._fiberinput:
            self._fiberinput['QSO'] = self._fiberinput['STAR']

    @property
    def fiberarea(self):
        """Average fiber area [arcsec^2] used for fiber input calculations"""
        return np.pi * self.fiberdia**2 / 4.0

    def extinction(self, wavelength):
        """
        Return atmospheric extinction [magnitudes/airmass]
        evaluated at wavelength (float or array)
        """
        return np.interp(wavelength, self._wave, self._extinction)

    def atmospheric_throughput(self, wavelength, airmass=1.0):
        """
        Return atmospheric throughput [0 - 1] at given airmass,
        evaluated at wavelength (float or array)
        """
        ext = self.extinction(wavelength)
        return 10**(-0.4 * airmass * ext)

    def fiberinput_throughput(self, wavelength=None, objtype=ObjType.STAR):
        """
        Return fiber input geometric throughput [0 - 1]
        evaluated at wavelength (float or array).
        If wavelength is None, do not interpolate.
        """
        if objtype in self._fiberinput:
            t = self._fiberinput[objtype]
        else:
            msg = 'Unknown objtype {}; using default fiber input loss'.format(objtype)
            msg += '\nKnown objtypes are '+str(list(self._fiberinput.keys()))
            warnings.warn(msg)
            t = self._fiberinput['default']

        if wavelength is None:
            return t
        else:
            return np.interp(wavelength, self._wave, t)

    def hardware_throughput(self, wavelength):
        """
        Return hardware throughput (optics, fiber run, CCD, but
        not including atmosphere or geometric fiber input losses)
        evaluated at wavelength (float or array)
        """
        return np.interp(wavelength, self._wave, self._thru, left=0.0, right=0.0)

    def _throughput(self, objtype=ObjType.STAR, airmass=1.0):
        """
        Returns system throughput for this object type and airmass
        at the native wavelengths of this Throughput object (self._wave)

        objtype may be any of the ObjType enumerated types
            CALIB : atmospheric extinction and fiber input losses not applied
            SKY   : fiber input losses are not applied
            other : all throughput losses are applied
        """
        objtype = objtype.strip().upper()

        Tatm = 10**(-0.4*airmass*self._extinction)
        if objtype == ObjType.CALIB:
            T = self._thru
        elif objtype == ObjType.SKY:
            T = self._thru * Tatm
        else:
            Tfiber = self.fiberinput_throughput(wavelength=None, objtype=objtype)
            T = self._thru * Tatm * Tfiber

        return T

    def __call__(self, wavelength, objtype=ObjType.STAR, airmass=1.0):
        """
        Returns system throughput at requested wavelength(s)

        objtype may be any of the ObjType enumerated types
            CALIB : atmospheric extinction and fiber input losses not applied
            SKY   : fiber input losses are not applied
            other : all throughput losses are applied
        """

        T = self._throughput(objtype=objtype, airmass=airmass)
        return np.interp(wavelength, self._wave, T, left=0.0, right=0.0)

    def thru(self, *args, **kwargs):
        """
        same as calling self(*args, **kwargs)
        """
        return self(*args, **kwargs)

    def photons(self, wavelength, flux, units="erg/s/cm^2/A", \
                objtype="STAR", exptime=None, airmass=1.0):
        """
        Returns photons per bin given input flux vs. wavelength,
        flux units, object type, exposure time, and airmass.

        NOTE: Returns raw photons per wavelength bin,
              *not* photons/A sampled at these wavelengths

        Inputs
        ------
        wavelength : input wavelength array in Angstroms
        flux       : input flux; same length as `wavelength`
        units      : units of `flux`
          * Treated as delta functions at each given wavelength:
            - "photons"
            - "erg/s/cm^2"
          * Treated as function values to be multipled by bin width:
            - "photon/A"
            - "erg/s/cm^2/A"
            - "erg/s/cm^2/A/arcsec^2"

        "photon" and "photon/A" are as observed by CCD and are returned
        without applying throughput terms.

        Optional Inputs
        ---------------
        objtype : string, optional; object type for Throughput object.
            SKY - atmospheric extinction and telescope+instrument throughput
                    applied, but not fiber input geometric losses.
            CALIB - telescope+instrument throughtput applied, but not
                    atmospheric extinction or fiber input geometric losses.
            Anything else (default) - treated as astronomical object
                    with all throughput terms applied.
        exptime : float, optional; exposure time, default self.exptime
        airmass : float, optional, default 1.0

        Returns
        -------
        array of number of photons observed by CCD at each wavelength,
        i.e. not per-Angstrom.

        Stephen Bailey, LBL
        May 2013
        """

        #- Wavelength bin size
        dw = np.gradient(wavelength)

        #- Standardize units; allow some sloppiness
        units = units.strip()  #- FITS pads short strings with spaces (!)
        units = units.replace("ergs", "erg")
        units = units.replace("photons", "photon")
        units = units.replace("Angstroms", "A")
        units = units.replace("Angstrom", "A")
        units = units.replace("Ang", "A")
        units = units.replace("**", "^")
        units = units.replace("cm2", "cm^2")
        units = units.replace("arcsec2", "arcsec^2")

        #- Check for units prefactor like "1e-17 erg/s/cm^2/A"
        scale = 1.0
        tmp = units.split()
        if len(tmp) == 2:
            try:
                scale = float(tmp[0])
                flux = flux * scale
                units = tmp[1]
            except ValueError:
                raise ValueError("Non-numeric units scale factor {}".format(tmp[0]))

        #- Default exposure time
        if exptime is None:
            exptime = self.exptime

        #- Input photons; return photons per bin (not photons per Angstrom)
        if units == "photon":
            return flux
        elif units == "photon/A":
            return flux * dw

        #- Sanity check on units
        if not units.startswith('erg'):
            raise ValueError("Unrecognized units {}".format(units))

        #- If we got here, we need to apply throughputs
        flux = self.apply_throughput(wavelength, flux,
                                 objtype=objtype, airmass=airmass)

        #- Convert to photons
        phot = flux * wavelength / self._hc

        #- erg/s/cm^2/A (i.e. Flambda, astronomical object)
        if units == "erg/s/cm^2/A":
            return phot * exptime * self.area * dw

        #- erg/s/cm^2/A/arcsec^2 (e.g. sky)
        elif units == "erg/s/cm^2/A/arcsec^2":
            return phot * exptime * self.area * dw * self.fiberarea

        #- erg/s/cm^2 (not per A; flux delta functions at given wavelengths)
        elif units == "erg/s/cm^2":
            return phot * exptime * self.area

        #- erg/s/cm^2/arcsec^2 (not per A; intensity delta functions)
        elif units == "erg/s/cm^2/arcsec^2":
            return phot * exptime * self.area * self.fiberarea

        else:
            raise ValueError("Unrecognized units {}".format(units))

    def apply_throughput(self, wavelength, flux, objtype="STAR", airmass=1.0):
        """
        Returns flux array with throughputs applied for given
        objtype and airmass.

        TODO: this is a simple throughput model that can be wrong if there
        is meaningful structure smaller than the wavelength sampling.
        """

        if flux.ndim == 1 or isinstance(objtype, str):
            thru = self.thru(wavelength, objtype=objtype, airmass=airmass)
            return flux * thru
        else:
            assert flux.ndim == 2
            assert flux.shape[0] == len(objtype)

            t = dict()
            objtype = np.array(objtype)
            outflux = flux.copy()
            for xt in set(objtype):
                thru = self.thru(wavelength, objtype=xt, airmass=airmass)
                ii = np.where(objtype == xt)[0]
                outflux[ii] *= thru

            return outflux

    @property
    def wavemin(self):
        """Minimum wavelength [Angstroms] covered by this throughput model"""
        return self._wave[0]

    @property
    def wavemax(self):
        """Maximum wavelength [Angstroms] covered by this throughput model"""
        return self._wave[-1]


# Copyright 2021 NREL

# 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.

import numpy as np
import pandas as pd

from floris.utilities import wrap_360
from pandas.core.base import DataError

from ..dataframe_operations import dataframe_manipulations as dfm
from ..energy_ratio import energy_ratio_visualization as ervis

class energy_ratio:
    """This class is used to calculate the energy ratios for a single
    dataframe with measurements, either from FLORIS or from SCADA data.
    This class supports bootstrapping for uncertainty quantification,
    automatic derivation of the frequency of bins based on occurrence
    in the provided dataset, and various choices for binning and daa
    discretization.
    """

    def __init__(self, df_in, inflow_freq_interpolant=None, verbose=False):
        """Initialization of the class.

        Args:
            df_in ([pd.DataFrame]): The dataframe provided by the user. This
            dataframe should have the following columns:
                * Reference wind direction for the test turbine, 'wd'
                * Reference wind speed for the test turbine, 'ws'
                * Power production of every turbine: pow_000, pow_001, ...
                * Reference power production used to normalize the energy
                    ratio: 'pow_ref'
            inflow_freq_interpolant (interpolant, optional): This is an
            interpolant that takes as inputs the wind direction and wind
            speed, and then returns the frequency of occurrence for that set
            of inflow conditions. If None is specified, the occurrence of each
            bin is derived from the provided data, df_in. Defaults to None.
            verbose (bool, optional): Print to console. Defaults to False.
        """
        self.verbose = verbose

        # Initialize dataframe
        self._set_df(df_in)

        # Initialize frequency functions
        self._set_inflow_freq_interpolant(inflow_freq_interpolant)

    # Private methods

    def _set_inflow_freq_interpolant(self, inflow_freq_interpolant):
        self.inflow_freq_interpolant = inflow_freq_interpolant

    def _set_df(self, df_in):
        """This function writes the dataframe provided by the user to the
        class as self.df_full. This full dataframe will be used to create
        a minimal dataframe called self.df which contains the minimum
        columns required to calculate the energy ratios. The contents of
        self.df will depend on the test_turbine specified and hence
        that dataframe is created in the _set_test_turbines() function.

        Args:
            df_in ([pd.DataFrame]): The dataframe provided by the user. This
            dataframe should have the following columns:
                * Reference wind direction for the test turbine, 'wd'
                * Reference wind speed for the test turbine, 'ws'
                * Power production of every turbine: pow_000, pow_001, ...
                * Reference power production used to normalize the energy
                    ratio: 'pow_ref'
        """
        if "pow_ref" not in df_in.columns:
            raise KeyError("pow_ref column not in dataframe. Cannot proceed.")
            # INFO: You can add such a column using:
            #   from flasc.dataframe_operations import \
            #       dataframe_manipulations as dfm
            #
            #   df = dfm.set_pow_ref_by_*(df)
            #   ...

        # Copy full dataframe to self
        self.df_full = df_in.copy()  # Full dataframe
        self.df = None

    def _set_test_turbines(self, test_turbines):
        """This function calculates the power production upon which the
        energy ratio is calculated, in the nominator of the energy ratio
        equation. This is typically a single turbine, e.g.,
        test_turbines=[0], but can also be the average of multiple turbines,
        e.g., test_turbines=[0, 1, 2]. This function creates the minimal
        dataframe, self.df, with columns being the wind direction 'wd',
        the wind speed 'ws', the power production of the test turbine(s)
        'pow_test' and the reference power production 'pow_ref'. The
        arrays 'pow_test' and 'pow_ref' are in the nominator and
        denominator in the energy ratio equation, respectively.

        Args:
            test_turbines ([iteratible]): List with the test turbine(s)
                used to calculate the power production in the nominator of
                the energy ratio equation. Typically, this is a single
                turbine, e.g., test_turbines=[0], but can also be multiple
                turbines. If multiple turbines are specified, it averages
                the power production between the turbines to come up with
                the test power values.
        """
        if not (type(test_turbines) is list):
            test_turbines = [test_turbines]
        self.test_turbines = test_turbines

        if "ti" in self.df_full.columns:
            cols = ["wd", "ws", "ti", "pow_ref"]
        else:
            cols = ["wd", "ws", "pow_ref"]

        self.df = self.df_full[cols].copy()
        self.df["pow_test"] = dfm.get_column_mean(
            df=self.df_full,
            col_prefix="pow",
            turbine_list=self.test_turbines,
            circular_mean=False,
        )

    def _set_binning_properties(
        self, ws_step=None, wd_step=None, wd_bin_width=None,
        ws_bins=None, wd_bins=None,
    ):
        """This function prepares the wind direction and wind speed bins in
        accordance to the user specified functions. Previously, the user could
        only specify the bins by assigning a ws_step and wd_step. Now, you can
        also specify the bins by assigning them directly. If 'ws_bins' is
        provided, then the variable 'ws_step' is completely ignored and the
        wind speed bins are directly set as the user-provided values in
        'ws_bins'. The same holds for wd_bins. If 'wd_bins' is provided, the
        variables 'wd_step' and 'wd_bin_width' are ignored and the wind
        direction bins are directly assigned the values in 'wd_bins'.

        Args:
            ws_step (float): Wind speed bin width and defines the step size
                at which the energy ratios are calculated along the wind speed.
                If 'ws_bins' is also provided, this variable is ignored.
            wd_step (float): Wind direction bin width and defines the step
                size at which the energy ratios are calculated along the wind
                direction. If 'wd_bins' is also provided, this variable is
                ignored.
            wd_bin_width (float, optional): Width of each wind direction bin.
                If this is larger than wd_step, there is overlap in the energy
                ratios between bins. This means data points are used more than
                once -- i.e., fall into multiple bins at the same time. If None
                is specified, defaults to the same value as wd_step. Note that
                if 'wd_bins' is also provided, this variable is ignored. Defaults
                to None.
            ws_bins (array, optional): Array containing the bins over which
                the energy ratios must be calculated (wind speeds). Each entry
                of the provided array must contain exactly two float values,
                being the lower and upper bound for that wind speed bin.
                Overlap between bins is not supported for wind speed bins,
                currently. Defaults to None.
            wd_bins (array, optional): Array containing the bins over which
                the energy ratios must be calculated (wind dir.). Each entry
                of the provided array must contain exactly two float values,
                being the lower and upper bound for that wind dir. bin.
                Overlap between bins is supported for wind direction bins.
                Defaults to None.
        """
        if ((ws_bins is None) | (wd_bins is None)):
            # Add a temporary variable
            a = np.array([-0.5, 0.5], dtype=float)

        # If ws_bins is not specified, automatically calculate the
        # wind speed bins using ws_step as bin width, from ws_step / 2.0,
        # which is bounded by [0.0, ws_step),  up to 30 m/s. If the user
        # already provided the wind speed bins 'ws_bins', then we need not
        # derive anything and we directly save the user-specified bins
        # to self (after converting them to numpy arrays).
        if wd_bins is None:
            ws_step = float(ws_step)
            ws_labels = np.arange(ws_step/2.0, 30.0001, ws_step)
            ws_bins = np.array([ws + a * ws_step for ws in ws_labels])
        else:
            ws_labels = np.array([np.mean(b) for b in ws_bins], dtype=float)
            ws_bins = np.array(ws_bins, dtype=float)

        # If wd_bins is not specified, automatically calculate the wind
        # direction bins using wd_step as bin width, from 0.0 deg to 360 deg.
        # If the user has already provided the wind direction bins 'ws_bins',
        # then we need not derive anything and we directly save the user
        # specified bins to self (after converting them to numpy arrays).
        if wd_bins is None:
            wd_step = float(wd_step)
            if wd_bin_width is None:
                wd_bin_width = wd_step
            wd_bin_width = float(wd_bin_width)

            wd_min = np.min([wd_step / 2.0, wd_bin_width / 2.0])
            wd_labels = np.arange(wd_min, 360.0001, wd_step)
            wd_bins = np.array([wd + a * wd_bin_width for wd in wd_labels])
        else:
            wd_labels = np.array([np.mean(b) for b in wd_bins], dtype=float)
            wd_bins = np.array(wd_bins, dtype=float)

        # Save variables
        self.ws_step = ws_step
        self.wd_step = wd_step
        self.wd_bin_width = wd_bin_width
        self.ws_labels = ws_labels
        self.wd_labels = wd_labels
        self.ws_bins = ws_bins
        self.wd_bins = wd_bins

    def _calculate_bins(self):
        """This function bins the data in the minimal dataframe, self.df,
        into the respective wind direction and wind speed bins. Note that
        there might be bin overlap if the specified wd_bin_width is larger
        than the bin step size. This code will copy dataframe rows that fall
        into multiple bins, effectively increasing the sample size.
        """
        # Bin according to wind speed. Note that data never falls into
        # multiple wind speed bins at the same time.
        for ws_bin in self.ws_bins:
            ws_interval = pd.Interval(ws_bin[0], ws_bin[1], "left")
            ids = (self.df["ws"] >= ws_bin[0]) & (self.df["ws"] < ws_bin[1])
            self.df.loc[ids, "ws_bin"] = np.mean(ws_bin)
            self.df.loc[ids, "ws_bin_edges"] = ws_interval

        # Bin according to wind direction. Note that data can fall into
        # multiple wind direction bins at the same time, if wd_bin_width is
        # larger than the wind direction binning step size, wd_step. If so,
        # data will be copied and the sample size is effectively increased
        # so that every relevant bin has that particular measurement.
        df_list = [None for _ in range(len(self.wd_labels))]
        for ii, wd_bin in enumerate(self.wd_bins):
            wd_interval = pd.Interval(wd_bin[0], wd_bin[1], "left")
            lb = wrap_360(wd_bin[0])
            ub = wrap_360(wd_bin[1])
            if ub < lb:  # Deal with angle wrapping
                ids = (self.df["wd"] >= lb) | (self.df["wd"] < ub)
            else:
                ids = (self.df["wd"] >= lb) & (self.df["wd"] < ub)
            df_subset = self.df.loc[ids].copy()
            df_subset["wd_bin"] = np.mean(wd_bin)
            df_subset["wd_bin_edges"] = wd_interval
            df_list[ii] = df_subset
        self.df = pd.concat(df_list, copy=False)

        # Make sure a float
        self.df["ws_bin"] = self.df["ws_bin"].astype(float)
        self.df["wd_bin"] = self.df["wd_bin"].astype(float)

    def _get_df_freq(self):
        """This function derives the frequency of occurrence of each bin
        (wind direction and wind speed) from the binned dataframe. The
        found values are used in the energy ratio equation to weigh the
        power productions of each bin according to their frequency of
        occurrence.
        """
        # Determine observed frequency
        cols = ["ws_bin", "wd_bin", "ws_bin_edges", "wd_bin_edges"]
        df_freq_observed = self.df[cols].copy()
        df_freq_observed["freq"] = 1
        df_freq_observed = df_freq_observed.groupby(["wd_bin", "ws_bin"])
        bin_edges = df_freq_observed[["ws_bin_edges", "wd_bin_edges"]].first()
        bin_freq = df_freq_observed["freq"].sum()
        df_freq_observed = pd.concat([bin_freq, bin_edges], axis=1)
        df_freq_observed = df_freq_observed.reset_index(drop=False)
        df_freq = df_freq_observed

        if self.inflow_freq_interpolant is not None:
            # Overwrite freq of bin occurrence with user-specified function
            df_freq["freq"] = self.inflow_freq_interpolant(
                df_freq["wd_bin"],
                df_freq["ws_bin"],
            )

        # Sort by 'ws_bin' as index
        df_freq = df_freq.set_index("ws_bin")
        self.df_freq = df_freq

        return df_freq

    # Public methods

    def get_energy_ratio(
        self,
        test_turbines,
        wd_step=2.0,
        ws_step=1.0,
        wd_bin_width=None,
        wd_bins=None,
        ws_bins=None,
        N=1,
        percentiles=[5.0, 95.0],
        return_detailed_output=False,
    ):
        """This is the main function used to calculate the energy ratios
        for dataframe provided to the class during initialization. One
        can calculate the energy ratio for different (sets of) turbines
        and under various discretization options.

        Args:
            test_turbines ([iteratible]): List with the test turbine(s)
                used to calculate the power production in the nominator of
                the energy ratio equation. Typically, this is a single
                turbine, e.g., test_turbines=[0], but can also be multiple
                turbines. If multiple turbines are specified, it averages
                the power production between the turbines to come up with
                the test power values.
            wd_step (float, optional): Wind direction discretization step
                size. This defines for what wind directions the energy ratio
                is to be calculated. Note that this does not necessarily
                also mean each bin has a width of this value. Namely, the
                bin width can be specified separately. Defaults to 2.0.
            ws_step (float, optional): Wind speed discretization step size.
                This defines the resolution and widths of the wind speed
                bins. Defaults to 1.0.
            wd_bin_width ([type], optional): The wind direction bin width.
                This value should be equal or larger than wd_step. When no
                value is specified, will default to wd_bin_width = wd_step.
                In the literature, it is not uncommon to specify a bin width
                larger than the step size to cover for variability in the
                wind direction measurements. By setting a large value for
                wd_bin_width, one gets a better idea of the larger-scale
                wake losses in the wind farm. Defaults to None.
            ws_bins (array, optional): Array containing the bins over which
                the energy ratios must be calculated (wind speeds). Each entry
                of the provided array must contain exactly two float values,
                being the lower and upper bound for that wind speed bin.
                Overlap between bins is not supported for wind speed bins,
                currently. Defaults to None.
            wd_bins (array, optional): Array containing the bins over which
                the energy ratios must be calculated (wind dir.). Each entry
                of the provided array must contain exactly two float values,
                being the lower and upper bound for that wind dir. bin.
                Overlap between bins is supported for wind direction bins.
                Defaults to None.
            N (int, optional): Number of bootstrap evaluations for
                uncertainty quantification (UQ). If N=1, will not perform
                any uncertainty quantification. Defaults to 1.
            percentiles (list, optional): Confidence bounds for the
                uncertainty quantification in percents. This value is only
                relevant if N > 1 is specified. Defaults to [5., 95.].
            return_detailed_output (bool, optional): Also calculate and
                return detailed energy ratio information useful for debugging
                and figuring out flaws in the data. This slows down the
                calculations but can be very useful. The additional info is
                written to self.df_lists[i]["er_results_info_dict"]. The
                dictionary variable therein contains two fields, being
                "df_per_wd_bin" and "df_per_ws_bin". The first gives an
                overview of the energy ratio for every wind direction bin,
                covering the collective effect of all wind speeds in the
                data. The latter one, "df_per_ws_bin", yields even more
                information and displays the energy ratio for every wind
                direction and wind speed bin, among others. This is
                particularly helpful in figuring out if the bins are well
                balanced. Defaults to False.

        Returns:
            energy_ratios ([pd.DataFrame]): Dataframe containing the found
                energy ratios under the prespecified settings. The dataframe
                contains the columns:
                    * wd_bin: The mean wind direction for this bin
                    * N_bin: Number of data entries in this bin
                    * baseline: Nominal energy ratio value (without UQ)
                    * baseline_l: Lower bound for energy ratio. This
                        value is equal to baseline without UQ and lower
                        with UQ.
                    * baseline_u: Upper bound for energy ratio. This
                        value is equal to baseline without UQ and higher
                        with UQ.
        """
        if self.df_full.shape[0] < 1:
            # Empty dataframe, do nothing
            self.energy_ratio_out = pd.DataFrame()
            self.energy_ratio_N = N
            return None

        if self.verbose:
            print("Calculating energy ratios with N = %d." % N)

        # Set up a 'pow_test' column in the dataframe
        self._set_test_turbines(test_turbines)

        # Set up bins
        self._set_binning_properties(
            ws_step=ws_step, wd_step=wd_step, wd_bin_width=wd_bin_width,
            ws_bins=ws_bins, wd_bins=wd_bins
        )
        self._calculate_bins()

        # Get probability distribution of bins
        self._get_df_freq()

        # Calculate the energy ratio for all bins
        out = _get_energy_ratios_all_wd_bins_bootstrapping(
            df_binned=self.df,
            df_freq=self.df_freq,
            N=N,
            percentiles=percentiles,
            return_detailed_output=return_detailed_output,
        )
        if return_detailed_output:
            energy_ratios = out[0]
            dict_out = out[1]
        else:
            energy_ratios = out

        self.energy_ratio_out = energy_ratios
        self.energy_ratio_N = N

        if return_detailed_output:
            return energy_ratios, dict_out

        return energy_ratios

    def get_energy_ratio_fast(
        self, test_turbines, ws_step, wd_step, wd_bin_width=None,
        ws_bins=None, wd_bins=None,
    ):
        """This function calculates the energy ratio in a fast manner
        but completely ignores any frequency weighing to achieve these speed-
        ups. It also does not support bootstrapping because of this.

        Args:
            test_turbines ([iteratible]): List with the test turbine(s)
                used to calculate the power production in the nominator of
                the energy ratio equation. Typically, this is a single
                turbine, e.g., test_turbines=[0], but can also be multiple
                turbines. If multiple turbines are specified, it averages
                the power production between the turbines to come up with
                the test power values.
            wd_step (float, optional): Wind direction discretization step
                size. This defines for what wind directions the energy ratio
                is to be calculated. Note that this does not necessarily
                also mean each bin has a width of this value. Namely, the
                bin width can be specified separately. Defaults to 2.0.
            ws_step (float, optional): Wind speed discretization step size.
                This defines the resolution and widths of the wind speed
                bins. Defaults to 1.0.
            wd_bin_width ([type], optional): The wind direction bin width.
                This value should be equal or larger than wd_step. When no
                value is specified, will default to wd_bin_width = wd_step.
                In the literature, it is not uncommon to specify a bin width
                larger than the step size to cover for variability in the
                wind direction measurements. By setting a large value for
                wd_bin_width, one gets a better idea of the larger-scale
                wake losses in the wind farm. Defaults to None.

        Returns:
            [type]: [description]
        """
        # Choose default option
        if wd_bin_width is None:
            wd_bin_width = wd_step

        # Set up a 'pow_test' column in the dataframe
        self._set_test_turbines(test_turbines)

        # Set up bins
        self._set_binning_properties(
            ws_step, wd_step, wd_bin_width, ws_bins, wd_bins
        )
        self._calculate_bins()

        df = self.df.dropna(how="any", subset=["pow_ref", "pow_test"]).copy()
        df.loc[df.index, "bin_count"] = 1

        df_summed = df.groupby("wd_bin")[["pow_ref", "pow_test", "bin_count"]].sum()
        df_summed.loc[df_summed.index, ["baseline", "baseline_lb", "baseline_ub"]] = (
            np.tile(df_summed["pow_test"] / df_summed["pow_ref"], (3, 1)).T
        )
        energy_ratios = df_summed.reset_index(drop=False)
        return energy_ratios

    def plot_energy_ratio(self):
        """This function plots the energy ratio against the wind direction,
        potentially with uncertainty bounds if N > 1 was specified by
        the user. One must first run get_energy_ratio() before attempting
        to plot the energy ratios.

        Returns:
            ax [plt.Axes]: Axis handle for the figure.
        """
        return ervis.plot(self.energy_ratio_out)

# Support functions not included in energy_ratio class

def _get_energy_ratios_all_wd_bins_bootstrapping(
    df_binned,
    df_freq,
    N=1,
    percentiles=[5.0, 95.0],
    return_detailed_output=False,
):
    """Wrapper function that calculates the energy ratio for every wind
    direction bin in the provided dataframe. This function wraps around
    the function '_get_energy_ratio_single_wd_bin_bootstrapping', which
    calculates the energy ratio for a single wind direction bin.

    Args:
        df_binned ([pd.DataFrame]): Dataframe containing the binned
        data. This dataframe must contain, at the minimum, the following
        columns:
            * ws_bin: The wind speed bin
            * wd_bin: The wind direction bin
            * pow_ref: The reference power production, previously specified
                by the user outside of this function/class. This value
                belongs in the denominator in the energy ratio equation.
            * pow_test: The test power production. This value belongs in the
                nominator in the energy ratio equation.
        df_freq ([pd.DataFrame]): Dataframe containing the frequency of every
            wind direction and wind speed bin. This dataframe is typically
            derived from the data itself but can also be a separate dataframe
            based on the wind rose of the site.
        N (int, optional): Number of bootstrap evaluations for
            uncertainty quantification (UQ). If N=1, will not perform any
            uncertainty quantification. Defaults to 1.
        percentiles (list, optional): Confidence bounds for the
            uncertainty quantification in percents. This value is only
            relevant if N > 1 is specified. Defaults to [5., 95.].
        return_detailed_output (bool, optional): Also calculate and
            return detailed energy ratio information useful for debugging
            and figuring out flaws in the data. This slows down the
            calculations but can be very useful. The additional info is
            written to self.df_lists[i]["er_results_info_dict"]. The
            dictionary variable therein contains two fields, being
            "df_per_wd_bin" and "df_per_ws_bin". The first gives an
            overview of the energy ratio for every wind direction bin,
            covering the collective effect of all wind speeds in the
            data. The latter one, "df_per_ws_bin", yields even more
            information and displays the energy ratio for every wind
            direction and wind speed bin, among others. This is
            particularly helpful in figuring out if the bins are well
            balanced. Defaults to False.

    Returns:
        energy_ratios ([pd.DataFrame]): Dataframe containing the found
            energy ratios under the prespecified settings. The dataframe
            contains the columns:
                * wd_bin: The mean wind direction for this bin
                * N_bin: Number of data entries in this bin
                * baseline: Nominal energy ratio value (without UQ)
                * baseline_l: Lower bound for energy ratio. This
                    value is equal to baseline without UQ and lower
                    with UQ.
                * baseline_u: Upper bound for energy ratio. This
                    value is equal to baseline without UQ and higher
                    with UQ.
    """
    # Extract minimal dataframe
    if "ti" in df_binned.columns:
        min_cols = [
            "wd",
            "ws",
            "ti",
            "ws_bin",
            "wd_bin",
            "pow_ref",
            "pow_test",
        ]
    else:
        min_cols = ["wd", "ws", "ws_bin", "wd_bin", "pow_ref", "pow_test"]
    df = df_binned[min_cols]

    # Save some relevant info
    unique_wd_bins = np.unique(df.wd_bin)
    # unique_ws_bins = np.unique(df.ws_bin)

    # Now calculate the actual energy ratios
    result = np.zeros([len(unique_wd_bins), 3])
    dict_out_list = [None for _ in range(len(unique_wd_bins))]

    for wd_idx, wd in enumerate(unique_wd_bins):
        df_subset = df[df["wd_bin"] == wd]
        df_freq_subset = df_freq[df_freq["wd_bin"] == wd]

        out = _get_energy_ratio_single_wd_bin_bootstrapping(
                df_binned=df_subset,
                df_freq=df_freq_subset,
                N=N,
                percentiles=percentiles,
                return_detailed_output=return_detailed_output,
        )
        if return_detailed_output:
            result[wd_idx, :] = out[0]
            dict_out_list[wd_idx] = out[1]
        else:
            result[wd_idx, :] = out

    # Save energy ratios to the dataframe
    df_out = pd.DataFrame(
        result, columns=["baseline", "baseline_lb", "baseline_ub"]
    )

    # Save wind direction bins and bin count to dataframe
    df_out["wd_bin"] = unique_wd_bins
    _, df_out["bin_count"] = np.unique(df["wd_bin"], return_counts=True)
    df_out["bin_count"] = df_out["bin_count"].astype(int)

    if return_detailed_output:
        # Concatenate dataframes and produce a new dict_out
        df_per_wd_bin = pd.concat([d["df_per_wd_bin"] for d in dict_out_list])
        df_per_ws_bin = pd.concat([d["df_per_ws_bin"] for d in dict_out_list])
        df_per_ws_bin = df_per_ws_bin.reset_index(drop=False)
        df_per_ws_bin = df_per_ws_bin.set_index(["wd_bin"])
        dict_out = {
            "df_per_wd_bin": df_per_wd_bin,
            "df_per_ws_bin": df_per_ws_bin,
        }
        return df_out, dict_out

    return df_out

def _get_energy_ratio_single_wd_bin_bootstrapping(
    df_binned,
    df_freq,
    N=1,
    percentiles=[5.0, 95.0],
    return_detailed_output=False,
):
    """Get the energy ratio for one particular wind direction bin and
    an array of wind speed bins. This function also includes bootstrapping
    functionality by increasing the number of bootstrap evaluations (N) to
    larger than 1. The bootstrap percentiles default to 5 % and 95 %.
    """
    # Get results excluding uncertainty
    if return_detailed_output:
        energy_ratio_nominal, dict_info = _get_energy_ratio_single_wd_bin_nominal(
            df_binned=df_binned,
            df_freq=df_freq,
            return_detailed_output=return_detailed_output,
        )
    else:
        energy_ratio_nominal = _get_energy_ratio_single_wd_bin_nominal(
            df_binned=df_binned,
            df_freq=df_freq,
            return_detailed_output=return_detailed_output,
        )

    # Add bootstrapping results, if necessary
    if N <= 1:
        results_array = np.array([energy_ratio_nominal] * 3, dtype=float)
    else:
        # Get a bootstrap sample of range
        bootstrap_results = np.zeros(N)
        bootstrap_results[0] = energy_ratio_nominal
        for i in range(1, N):
            df_randomized = df_binned.sample(frac=1, replace=True).copy()
            bootstrap_results[i] = _get_energy_ratio_single_wd_bin_nominal(
                df_binned=df_randomized,
                df_freq=df_freq,
                return_detailed_output=False,
            )

        # Return the results in the order used in previous versions
        results_array = np.array(
            [
                energy_ratio_nominal,
                np.nanpercentile(bootstrap_results, percentiles)[0],
                np.nanpercentile(bootstrap_results, percentiles)[1],
            ]
        )

    if return_detailed_output:
        return results_array, dict_info
    else:
        return results_array

def _get_energy_ratio_single_wd_bin_nominal(
    df_binned, df_freq=None, return_detailed_output=False
):
    """Get the energy ratio for one particular wind direction bin and
    an array of wind speed bins. This function performs a single
    calculation of the energy ratios without uncertainty quantification.
    """
    # Copy minimal dataframe
    if "ti" in df_binned.columns:
        min_cols = [
            "wd_bin",
            "ws_bin",
            "wd",
            "ws",
            "ti",
            "pow_ref",
            "pow_test",
        ]
        mean_cols = ["wd", "ws", "ti", "pow_ref", "pow_test"]
        std_cols = ["wd", "ws", "ti", "pow_ref", "pow_test"]

    else:
        min_cols = ["wd", "ws", "wd_bin", "ws_bin", "pow_ref", "pow_test"]
        mean_cols = ["wd", "ws", "pow_ref", "pow_test"]
        std_cols = ["wd", "ws", "pow_ref", "pow_test"]
    df = df_binned[min_cols].copy()

    # Drop any faulty measurements
    df = df.dropna(how="any")

    # Check if only one wd_bin present in data
    wd_bin = df_binned["wd_bin"].unique()
    if len(wd_bin) > 1:
        raise DataError("More than one wd_bin present in data.")

    # Reference and test turbine energy
    df["freq"] = 1
    df_sums = df.groupby("ws_bin")[["pow_ref", "pow_test", "freq"]].sum()
    df_sums.columns = [
        "energy_ref_unbalanced",
        "energy_test_unbalanced",
        "bin_count",
    ]

    if return_detailed_output:
        # Calculate bin information
        df_stds = df.groupby("ws_bin")[std_cols].std()
        df_stds.columns = ["{}_std".format(c) for c in df_stds.columns]

        # Mean values of bins and power values
        df_means = df.groupby("ws_bin")[mean_cols].mean()
        df_means.columns = ["{}_mean".format(c) for c in df_means.columns]

        # Collect into a single dataframe
        df_per_ws_bin = pd.concat([df_means, df_stds, df_sums], axis=1)
        df_per_ws_bin["wd_bin"] = wd_bin[0]
        df_per_ws_bin["wd_bin_edges"] = df_freq["wd_bin_edges"]
        df_per_ws_bin["ws_bin_edges"] = df_freq["ws_bin_edges"]

        # Calculate unbalanced energy ratio for each wind speed bin
        df_per_ws_bin["energy_ratio_unbalanced"] = (
            df_per_ws_bin["energy_test_unbalanced"]
            / df_per_ws_bin["energy_ref_unbalanced"]
        )

        # Calculate (total) unbalanced energy ratio for all wind speeds
        energy_ratio_total_unbalanced = (
            df_per_ws_bin["energy_test_unbalanced"].sum()
            / df_per_ws_bin["energy_ref_unbalanced"].sum()
        )

        # Calculate total statistics
        total_means = df[mean_cols].mean()
        total_means = total_means.rename(
            dict(zip(mean_cols, ["{:s}_mean".format(c) for c in mean_cols]))
        )
        total_stds = df[std_cols].std()
        total_stds = total_stds.rename(
            dict(zip(std_cols, ["{:s}_std".format(c) for c in std_cols]))
        )

        # Get summation of energy and bin frequencies
        total_sums = df[["pow_ref", "pow_test", "freq"]].sum()
        total_sums = total_sums.rename(
            {
                "pow_ref": "energy_ref_unbalanced",
                "pow_test": "energy_test_unbalanced",
                "freq": "bin_count",
            }
        )

        df_per_wd_bin = pd.concat([total_means, total_stds, total_sums])
        df_per_wd_bin["wd_bin"] = wd_bin[0]
        df_per_wd_bin[
            "energy_ratio_unbalanced"
        ] = energy_ratio_total_unbalanced

    else:
        df_per_wd_bin = pd.DataFrame({"wd_bin": [wd_bin[0]]})
        df_per_ws_bin = pd.DataFrame(
            {
                "pow_ref_mean": (
                    df_sums["energy_ref_unbalanced"] / df_sums["bin_count"]
                ),
                "pow_test_mean": (
                    df_sums["energy_test_unbalanced"] / df_sums["bin_count"]
                ),
                "bin_count": df_sums["bin_count"],
            }
        )

    # Write bin frequencies to the dataframe and ensure normalization
    df_per_ws_bin["freq_balanced"] = df_freq["freq"] / df_freq["freq"].sum()
    df_per_ws_bin["freq_balanced"] = df_per_ws_bin["freq_balanced"].fillna(0)

    # Calculate normalized balanced energy for ref and test turbine
    df_per_ws_bin["energy_ref_balanced_norm"] = (
        df_per_ws_bin["pow_ref_mean"] * df_per_ws_bin["freq_balanced"]
    )
    df_per_ws_bin["energy_test_balanced_norm"] = (
        df_per_ws_bin["pow_test_mean"] * df_per_ws_bin["freq_balanced"]
    )

    # Compute total balanced energy ratio over all wind speeds
    df_per_wd_bin["energy_test_balanced_norm"] = (
        df_per_ws_bin["energy_test_balanced_norm"].sum()
    )
    df_per_wd_bin["energy_ref_balanced_norm"] = (
        df_per_ws_bin["energy_ref_balanced_norm"].sum()
    )

    energy_ratio_total_balanced = float(
        df_per_wd_bin["energy_test_balanced_norm"] /
        df_per_wd_bin["energy_ref_balanced_norm"]
    )

    if return_detailed_output:
        df_per_ws_bin["energy_ratio_balanced"] = (
            df_per_ws_bin["energy_test_balanced_norm"]
            / df_per_ws_bin["energy_ref_balanced_norm"]
        )
        df_per_wd_bin["energy_ratio_balanced"] = energy_ratio_total_balanced

        # Formatting
        df_per_wd_bin = pd.DataFrame(df_per_wd_bin).T
        df_per_wd_bin["bin_count"] = df_per_wd_bin["bin_count"].astype(int)

        df_per_wd_bin["wd_bin_edges"] = df_freq.iloc[0]["wd_bin_edges"]
        df_per_wd_bin = df_per_wd_bin.set_index("wd_bin")

        df_per_ws_bin["bin_count"] = df_per_ws_bin["bin_count"].astype(int)
        dict_out = {
            "df_per_wd_bin": df_per_wd_bin,
            "df_per_ws_bin": df_per_ws_bin,
        }
        return energy_ratio_total_balanced, dict_out

    return energy_ratio_total_balanced