#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np


# In[2]:


def parallax_bias(ruwe, parallax, parallax_error, Gmag):
    '''
    This function takes arrays of ruwe, parallax, parallax_error, and Gmag and returns the expected true parallax error after uncertainty inflation.
    '''
    inflation_factor = f(ruwe = ruwe, parallax = parallax, Gmag = Gmag)
    return parallax_error*inflation_factor

def f(ruwe, parallax, Gmag, f0 = 3.73, alpha=2.77, beta=0.065,gamma=-0.056):
    '''
    This is Equation 3 of the submitted paper. 
    '''
    sigma_eta = al_uncertainty_per_ccd_interp(G = Gmag)
    fmax = f0*(parallax/10)**beta* (sigma_eta/0.1)**gamma
    fval = 1 + (fmax - 1)*(1 - np.exp(-alpha*(ruwe-1)))
    fval[fval < 1] = 1
    return fval

def al_uncertainty_per_ccd_interp(G):
    '''
    This gives the uncertainty *per CCD* (not per FOV transit), taken from Fig 3 of https://arxiv.org/abs/2206.05439
    This is the "EDR3 adjusted" line from that Figure, which is already inflated compared to the formal uncertainties.
    '''
    G_vals =    [ 4,    5,   6,     7,   8.2,  8.4, 10,    11,    12,  13,    14,   15,   16,   17,   18,   19,  20]
    sigma_eta = [0.4, 0.35, 0.15, 0.17, 0.23, 0.13,0.13, 0.135, 0.125, 0.13, 0.15, 0.23, 0.36, 0.63, 1.05, 2.05, 4.1]
    return np.interp(G, G_vals, sigma_eta)



# In[3]:


# example usage 
inflated_err = parallax_bias(ruwe = np.array([1.5]), parallax = np.array([2.0]), 
                             parallax_error = np.array([0.035]), Gmag = np.array([14]))

print(inflated_err)


# In[ ]: