import os
import simulators
import numpy as np
import json

"""Calculate Errors on the Spectrum.

   For a first go using an fixed SNR of 200 for all observations.
"""


def get_snrinfo(star, obs_num, chip):
    """Load SNR info from json file."""
    snr_file = os.path.join(simulators.paths["spectra"], "detector_snrs.json")
    with open(snr_file, "r") as f:
        snr_data = json.load(f)
    try:
        return snr_data[str(star)][str(obs_num)][str(chip)]
    except KeyError as e:
        print("No snr data present for {0}-{1}_{2}".format(star, obs_num, chip))
        raise e


def spectrum_error(star, obs_num, chip, error_off=False):
    """Return the spectrum error.

    errors = None will perform a normal chi**2 statistic.
    """
    if error_off:
        errors = None
    else:
        snr = get_snrinfo(star, obs_num, chip)
        if len(snr) == 1:
            errors = 1 / np.float(snr[0])
        else:
            raise NotImplementedError("Haven't checked if an error array can be handled yet.")
    return errors

import os
import simulators
import numpy as np
import json
import warnings

"""Calculate Errors on the Spectrum.

   For a first go using an fixed SNR of 200 for all observations.
"""


def get_snrinfo(star, obs_num, chip):
    """Load SNR info from json file."""
    snr_file = os.path.join(simulators.paths["spectra"], "detector_snrs.json")
    with open(snr_file, "r") as f:
        snr_data = json.load(f)
    try:
        return snr_data[str(star)][str(obs_num)][str(chip)]
    except KeyError as e:
        warnings.warn("No snr data present for {0}-{1}_{2}. "
                      "Setting error to None instead".format(star, obs_num, chip))
        return None


def spectrum_error(star, obs_num, chip, error_off=False):
    """Return the spectrum error.

    errors = None will perform a normal chi**2 statistic.
    """
    if error_off:
        errors = None
    else:
        snr = get_snrinfo(star, obs_num, chip)
        if snr is None:
            errors = None
        elif len(snr) == 1:
            errors = 1 / np.float(snr[0])
        else:
            raise NotImplementedError("Haven't checked if an error array can be handled yet.")
    return errors

Handle no snr information in snr file. (for fake simualtions mainly)
Handle no snr information in snr file. (for fake simualtions mainly)

  import os
  import simulators
  import numpy as np
  import json
+ import warnings
  
  """Calculate Errors on the Spectrum.
  
     For a first go using an fixed SNR of 200 for all observations.
  """
  
  
  def get_snrinfo(star, obs_num, chip):
      """Load SNR info from json file."""
      snr_file = os.path.join(simulators.paths["spectra"], "detector_snrs.json")
      with open(snr_file, "r") as f:
          snr_data = json.load(f)
      try:
          return snr_data[str(star)][str(obs_num)][str(chip)]
      except KeyError as e:
-         print("No snr data present for {0}-{1}_{2}".format(star, obs_num, chip))
-         raise e
+         warnings.warn("No snr data present for {0}-{1}_{2}. "
+                       "Setting error to None instead".format(star, obs_num, chip))
+         return None
  
  
  def spectrum_error(star, obs_num, chip, error_off=False):
      """Return the spectrum error.
  
      errors = None will perform a normal chi**2 statistic.
      """
      if error_off:
          errors = None
      else:
          snr = get_snrinfo(star, obs_num, chip)
+         if snr is None:
+             errors = None
-         if len(snr) == 1:
+         elif len(snr) == 1:
              errors = 1 / np.float(snr[0])
          else:
              raise NotImplementedError("Haven't checked if an error array can be handled yet.")
      return errors
  
Handle no snr information in snr file. (for fake simualtions mainly)
## Code Before:
import os
import simulators
import numpy as np
import json

"""Calculate Errors on the Spectrum.

   For a first go using an fixed SNR of 200 for all observations.
"""


def get_snrinfo(star, obs_num, chip):
    """Load SNR info from json file."""
    snr_file = os.path.join(simulators.paths["spectra"], "detector_snrs.json")
    with open(snr_file, "r") as f:
        snr_data = json.load(f)
    try:
        return snr_data[str(star)][str(obs_num)][str(chip)]
    except KeyError as e:
        print("No snr data present for {0}-{1}_{2}".format(star, obs_num, chip))
        raise e


def spectrum_error(star, obs_num, chip, error_off=False):
    """Return the spectrum error.

    errors = None will perform a normal chi**2 statistic.
    """
    if error_off:
        errors = None
    else:
        snr = get_snrinfo(star, obs_num, chip)
        if len(snr) == 1:
            errors = 1 / np.float(snr[0])
        else:
            raise NotImplementedError("Haven't checked if an error array can be handled yet.")
    return errors

## Instruction:
Handle no snr information in snr file. (for fake simualtions mainly)
## Code After:
import os
import simulators
import numpy as np
import json
import warnings

"""Calculate Errors on the Spectrum.

   For a first go using an fixed SNR of 200 for all observations.
"""


def get_snrinfo(star, obs_num, chip):
    """Load SNR info from json file."""
    snr_file = os.path.join(simulators.paths["spectra"], "detector_snrs.json")
    with open(snr_file, "r") as f:
        snr_data = json.load(f)
    try:
        return snr_data[str(star)][str(obs_num)][str(chip)]
    except KeyError as e:
        warnings.warn("No snr data present for {0}-{1}_{2}. "
                      "Setting error to None instead".format(star, obs_num, chip))
        return None


def spectrum_error(star, obs_num, chip, error_off=False):
    """Return the spectrum error.

    errors = None will perform a normal chi**2 statistic.
    """
    if error_off:
        errors = None
    else:
        snr = get_snrinfo(star, obs_num, chip)
        if snr is None:
            errors = None
        elif len(snr) == 1:
            errors = 1 / np.float(snr[0])
        else:
            raise NotImplementedError("Haven't checked if an error array can be handled yet.")
    return errors
