# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/sarahrdk/EscapeEarth/blob/main/Interns/Sarah/Activity_3.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="pb2ZyejuQpcY"
# # Mount google drive to access/write data

# + id="UKPF4ZnDQZ4t" colab={"base_uri": "https://localhost:8080/"} outputId="c4d617a7-6166-408a-db94-f1945b1dff1b"
from google.colab import drive

drive.mount('/content/gdrive')

# + [markdown] id="doUQI5KgQuhG"
# # Import/Install needed packages 

# + id="hSvwR_JyQa6X"
import pandas as pd
# !pip install lightkurve
import lightkurve as lk

# + [markdown] id="PIcVbkjfUVN9"
#
# # Functions to open data

# + id="f425gZFZUZC1"
import lightkurve as lk
import glob
import os
import pandas as pd

def locate_files(tic,path=None):
    '''
    ~ Locates TESS lightcurve files with filenames formatted from a mast bulk download.~
    REQUIRES: glob
    Args: 
        tic            -(int or str)TESS TIC ID
        path           -(str) path on computer to file(s) location
    Returns:
        list of path strings for all files found with specified tic
    '''
    if path == None: #if only need filename
        fullpath = glob.glob('*{}-*-s_lc.fits'.format(tic)) #to use wildcard*
    else: #user defined path to datafile on their computer
        pathstart = path   
        pathstart = str(pathstart) #make a string in case user forgets to but think that gives an err anyway
        pathend = pathstart +'*{}-*-s_lc.fits'.format(tic) #stitches path & filename
        fullpath= glob.glob(pathend) #to use wildcard* 
    return fullpath


def open_rawdata(fullpath,sector):
    '''
    ~ Opens raw data light curve file objects downloaded to our shared google drive folder~
    REQUIRES: lightkurve as lk
    Args:
        fullpath   -(str) list of path strings for all files found with specified tic
        sector    -(int) sector number for desired data
    Returns:
        lcfs      -(list) list of lightkurve 'lightcurvefile' class objects
    '''
    lcfs = []
    for file in fullpath:
        if len(file)==0:
            print('no files')
        else:
            try:
                lcfile = lk.open(file)
                mystring = str(type(lcfile))
                if mystring[34:-2] == 'TessLightCurveFile':
                    hdr = lcfile.get_header()
                    mysector = hdr['SECTOR']
                    if sector == mysector:
                        lcfs.append(lcfile)
                    else:
                        pass
                else:
                    pass
            except FileNotFoundError:
                pass
    return lcfs


# + [markdown] id="AfKTaV_MQ43e"
# # Load in starting data

# + id="9e0Yl_vTQk8R" colab={"base_uri": "https://localhost:8080/"} outputId="25e1f275-7ab3-447e-fb90-d3747a29bba7"
# tic ids for subsample
subsample_tics = [7582594, 7582633, 7620704, 7618785, 7584049]
# path to raw data
raw_data_path = '/content/gdrive/MyDrive/EscapeEarthData/Sector_14_rawdata/'

# open data
paths2files = locate_files(subsample_tics[1],raw_data_path) #locate filename given tic id 
subsample_data = open_rawdata(paths2files,sector=14) #open the lightcurvefile
print('Initial function output data type:',type(subsample_data))

# + colab={"base_uri": "https://localhost:8080/"} id="uU0uI7WmenGH" outputId="325f097e-bb30-4fe0-eb1f-bb85757a1240"
list1 = []
for count, i in enumerate(subsample_tics):
  print(i)
  print(count)
  print(raw_data_path)
  paths2files = locate_files(i,raw_data_path) #locate filename given tic id 
  print(paths2files)
  #saveguard against missing files  
  if len(paths2files)==0: #if no files exist go to next target in for loop    
    pass  
  else: 
    subsample_data = open_rawdata(paths2files,sector=14) #open the lightcurvefile
    print(subsample_data)
    newclass = subsample_data[0].PDCSAP_FLUX
    hello = newclass.remove_outliers(sigma = 3)
    print(type(hello))
    list1.append(hello)

# + [markdown] id="ROrQ6gq0iFRM"
# # Test on a single target

# + id="nyUL2FMsSUZl" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="3cec96f1-2233-46cc-d7f0-91ff977200e3"
# examine input data

print('Initial data type:',type(list1[0])) #check inital class object type to see allowed methods
for i in list1:
  i.scatter() #visualize original data

# + [markdown] id="YUteG9flcrZO"
# # Apply class method to clean single target

# + id="6g--XlcIQk-w" colab={"base_uri": "https://localhost:8080/"} outputId="c3739d7e-508f-4002-f91e-e23813f0cd80"
# change/select data type to necessary class object in order to use class method

pdcsap_data = subsample_data[0].PDCSAP_FLUX #change class object type
print('original data type:',type(subsample_data[0])) #print to see change in class object type
print('new data type:',type(pdcsap_data)) #verify class object change

# + id="x0vaAuWwQlBH"
# apply class method


bin_data = pdcsap_data.bin() #apply method
bin_data.plot() #visualize the effect of method
print('type:',type(bin_data)) #check for output data type

# + id="5dq9xPpxciy9"
# show how to use any helpful args

bin_data2 = pdcsap_data.bin(bins=27) #see how args change output data
bin_data2.plot() #visualize effect of arg change

# + id="-jvxH4sFQa9j" colab={"base_uri": "https://localhost:8080/"} outputId="13ef2279-2307-43a5-8619-048788f99b52"
# examine output data

print('Output data type:',type(bin_data2)) #NOTE: your output data type should be 'TessLightCurve' or 'LightCurve'

# + [markdown] id="a4wOVBGqiLlf"
# # Apply method iteratively to all targets

# + id="85mLzkHbiP8r"
for target in subsample_tics:
  #open data
  paths = locate_files(target,raw_data_path)
  original_data = open_rawdata(paths,sector=14)

  #saveguard against missing files
  if len(original_data)==0: #if no files exist go to next target in for loop
    pass
  else: #if files exist proceed

    #format data (class object type)
    new_data = original_data[0].PDCSAP_FLUX #index into orignal_data b/c its a list
    #verify class object change
    print('class object change:',type(original_data[0]),'changed to',type(new_data))

    #apply method
    cleaned_data = new_data.bin(bins=27)
    #verify cleaned data looks different (or save it)
    cleaned_data.plot()

# + id="mv3KVLSiQbAt"
# BONUS work: write this new data to a file so we can open cleaned data directly
# instead of opening raw data, cleaning it, then running analysis code on it
# Tip: Similar to above, test your code on one target first, then edit code to 
# work in a for loop


StarcoderdataJupyterScriptsDedupFiltered