# Setup: The Libraries
# This code installs the necessary libraries if not already available, clones a github repository being used to hold
# example galaxies, and imports what we will be using as data analysis tools.
# Note: Do not change this code.
!pip install astropy matplotlib numpy

import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import simps
from astropy.io import fits

from google.colab import drive
drive.mount('/content/drive')
%cd /content/drive/MyDrive/
!git clone https://github.com/h-roundcount/DJA-Tutorial-Example-Galaxies.git

#Don't change this code; it is simply setting up the example galaxies for you.

repo_path = '/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies'
data_path = os.path.join(repo_path, 'data')
example_files = ['/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/gds-barrufet-s67-v2_prism-clear_2198_6620.spec (1).fits', '/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/gds-barrufet-s67-v2_prism-clear_2198_8290.spec (2).fits', "/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/gds-barrufet-s67-v2_prism-clear_2198_8777.spec (2).fits",
                  "/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/glazebrook-v2_prism-clear_2565_12629.spec (1).fits", "/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/glazebrook-v2_prism-clear_2565_13079.spec.fits",
                  "/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/glazebrook-v2_prism-clear_2565_14552.spec (2) (1).fits",
                  "/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/glazebrook-v2_prism-clear_2565_4232.spec.fits"]
example_redshifts = [4.3563, 4.6563, 3.4713, 3.1962, 3.1979, 3.542, 3.6921]

import matplotlib.pyplot as plt
import numpy as np
from astropy.io import fits

prominent_features = {
    "Ca II H": 3968,
    "Ca II K": 3934,
    "H-Delta": 4102,
    "H-Gamma": 4340,
    "H-Beta": 4861,
    "O III": 5007,
    "Mg I": 5175,
    "Fe I": 5270,
    "Na I": 5893,
    "H-alpha + N II": 6563,
    "S II": 6716,
}

def process_fits_file(fits_file, redshift):
    with fits.open(fits_file) as hdul:
        data = hdul[1].data
        wavelength = data['wave'] * 1e4  # Convert micrometers to Angstroms - DO NOT EDIT THIS LINE
        flux = data['flux']

    # Here, we convert to rest frame to standardize our spectrums.
    lambda_rest = wavelength / (1 + redshift)
    return lambda_rest, flux

def plot_quiescent_spectra(fits_files, redshifts):
    """
    Plots quiescent galaxy spectra from all of the FITS files you add to the list (fits_files),
    sorted by redshifts which you add to the other list (redshifts). Make sure to keep the redshifts
    in the same order that your files are so they correspond correctly.
    """
    plt.figure(figsize=(10, 8))

    #Here, we plot the spectrums themselves; wavelength versus flux.
    #If the lines are too close together, change offset to be larger.
    for i, (fits_file, z) in enumerate(zip(fits_files, redshifts)):
        offset = 2 #This is where you can change the offset value if you need.
        lambda_rest, flux = process_fits_file(fits_file, z)
        offset = i * offset
        plt.plot(lambda_rest, flux + offset, label=f"z = {z:.2f}")

    # Here we add vertical lines at the locations of all of the significant wavelengths.
    for feature, wavelength in prominent_features.items():
        plt.axvline(x=wavelength, color='gray', linestyle='--', alpha=0.7)
        plt.text(wavelength, 2.85 * len(fits_files), feature, rotation=90, verticalalignment='bottom', fontsize=8, alpha=0.7)

    # Plot everything.
    plt.xlabel("Rest Wavelength (Angstroms)")
    plt.xlim(0, 8000)
    plt.ylabel("Flux + Offset")
    plt.title("Quiescent Galaxy Spectra Sorted by Redshift")
    plt.legend(loc='upper right', fontsize=8)
    plt.grid(False)
    plt.tight_layout()
    plt.show()

# Run the code!
fits_files = example_files #replace this with a list of your own files, or edit the example_files list which appeared under the "Downloaded Data" section.
redshifts = example_redshifts #replace this with a list of your own values which correspond with the files you downloaded; this information is the z-spec column on the DJA.
plot_quiescent_spectra(fits_files, redshifts)

import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits

def convert_to_angstroms(wavelength_micrometers):
    return wavelength_micrometers * 1e4

def calculate_d4000(wavelength, flux, redshift):
    """
    Calculate the D4000 index.
    """
    # First, we convert to rest frame.
    rest_wavelength = wavelength / (1 + redshift)

    # Here we create wavelength ranges for D4000 calculation in Angstroms so we do not have to parse ALL of the data
    blue_range = (3750, 3950)
    red_range = (4050, 4250)
    blue_mask = (rest_wavelength >= blue_range[0]) & (rest_wavelength <= blue_range[1])
    red_mask = (rest_wavelength >= red_range[0]) & (rest_wavelength <= red_range[1])

    # Calculate mean fluxes in each range
    if np.any(blue_mask) and np.any(red_mask):
        mean_flux_blue = np.mean(flux[blue_mask])
        mean_flux_red = np.mean(flux[red_mask])
        return mean_flux_red / mean_flux_blue
    else:
        raise ValueError("No valid indices found for D4000 calculation.")

def calculate_ew_hdelta(wavelength, flux, redshift):
    """
    Calculate the equivalent width of Hδ.

    This function measures the width of the absorption or emission feature in the spectrum
    relative to the surrounding continuum. It is a common way to quantify the
    strength of a spectral line.
    """
    # First, we convert to rest frame (Are you noticing a pattern with this?)
    rest_wavelength = wavelength / (1 + redshift)

    # Create our ranges.
    hdelta_range = (4083, 4122)
    continuum_range = (4041, 4079)


    hdelta_mask = (rest_wavelength >= hdelta_range[0]) & (rest_wavelength <= hdelta_range[1])
    continuum_mask = (rest_wavelength >= continuum_range[0]) & (rest_wavelength <= continuum_range[1])



    if np.any(hdelta_mask) and np.any(continuum_mask):
        # Calculate continuum flux and Hδ flux
        continuum_flux = np.mean(flux[continuum_mask])
        hdelta_flux = np.sum(np.abs(flux[hdelta_mask] - continuum_flux))
        hdelta_width = np.ptp(rest_wavelength[hdelta_mask])  # Peak-to-peak wavelength range
        print("Continuum Flux:", continuum_flux)
        print("Hδ Flux:", hdelta_flux)
        return hdelta_flux / continuum_flux * hdelta_width
    else:
        raise ValueError("No valid indices found for Hδ equivalent width calculation.")


def process_fits_file(fits_file, redshift=0):
    """
    This takes the files you have and makes them usable. It also makes our
    flux and wavelength lists!
    """
    # Open the FITS file and extract wavelength and flux
    with fits.open(fits_file) as hdul:
        data = hdul[1].data
        wavelength = data['wave']
        flux = data['flux']

    # Convert units to Angstroms
    wavelength_angstroms = convert_to_angstroms(wavelength)

    # Calculate D4000 index
    d4000 = calculate_d4000(wavelength_angstroms, flux, redshift)

    # Calculate Hδ equivalent width
    ew_hdelta = calculate_ew_hdelta(wavelength_angstroms, flux, redshift)

    return d4000, ew_hdelta

def plot_d4000_vs_ew(fits_files, redshifts):
    """
    Plots D4000 index against Hδ equivalent width for all of your different FITS files.
    """
    # Some error catching. If this error runs, it means you probably accidentally
    # added a copy of a galaxy or redshift, or accidentally forgot to include one.
    if len(fits_files) != len(redshifts):
        raise ValueError("The number of FITS files and redshifts must match.")

    d4000_values = []
    ew_hdelta_values = []

    for fits_file, redshift in zip(fits_files, redshifts):
        try:
            print(f"Processing {fits_file} with redshift {redshift}...")
            d4000, ew_hdelta = process_fits_file(fits_file, redshift)
            d4000_values.append(d4000)
            ew_hdelta_values.append(ew_hdelta)
        except ValueError as e:
            print(f"Skipping file {fits_file}: {e}")
        except Exception as e:
            print(f"Unexpected error with file {fits_file}: {e}")

    # Plot!
    plt.figure(figsize=(8, 6))
    plt.scatter(d4000_values, ew_hdelta_values, color='blue', edgecolor='k')
    plt.ylim(bottom=-0.1, top=2)
    plt.xlabel('D4000 Index', fontsize=14)
    plt.ylabel('EW(Hδ) [Angstroms]', fontsize=14)
    plt.title('D4000 vs. EW(Hδ)', fontsize=16)
    plt.grid(True)
    plt.show()

#Run that last huge section of code!

fits_files = example_files
redshifts = example_redshifts
plot_d4000_vs_ew(fits_files, redshifts)

# This is going to require some outside research from you. The DJA doesn't have a convenient way of extracting color information yet.
# Therefore, you are required to provide the redshift and magnitudes as input parameters yourself.
# You can do this by finding papers which discuss the galaxies you're analyzing; for example, by using the NASA ADS: https://ui.adsabs.harvard.edu/
# The script will only use the .fits file for wavelength extraction.

def rest_frame_color_color_plot(fits_file, redshift, u_mag, v_mag, j_mag):
    """
    Creates a rest-frame UVJ diagram (U-V vs. V-J). You will need to find literature
    that includes the u,v, and j values you need, because they are not avaialable on the DJA.
    """
    # Open the .fits file and extract wavelength
    with fits.open(fits_file) as hdul:
        data = hdul[1].data
        wavelength = data['wave']

    # Conert to rest frame
    def correct_for_redshift(mag, z):
        return mag - 2.5 * np.log10(1 + z)

    u_rest = correct_for_redshift(u_mag, redshift)
    v_rest = correct_for_redshift(v_mag, redshift)
    j_rest = correct_for_redshift(j_mag, redshift)

    v_j_color = np.linspace(0.5, 2, 100)  # V-J color range

    u_v_color = 0.88 * v_j_color + 0.49  # Diagonal line equation
    u_v_horizontal = 1.3  # Horizontal line: U-V = 1.3
    v_j_vertical = 1.6    # Vertical line: V-J = 1.6
    v_j_intersection = (1.3 - 0.49) / 0.88
    u_v_intersection = 0.88 * 1.6 + 0.49
    plt.plot(v_j_color[(v_j_color >= v_j_intersection) & (v_j_color <= 1.6)],
            u_v_color[(v_j_color >= v_j_intersection) & (v_j_color <= 1.6)], c='blue', label=r'Diagonal Segment', alpha=0.7)
    plt.plot(np.linspace(0, v_j_intersection, 10), np.ones(10) * u_v_horizontal, c='blue', label=r'Horizontal Segment', alpha=0.7)
    plt.plot(np.ones(10) * 1.6, np.linspace(u_v_intersection, 3, 10), c='blue', label=r'Vertical Segment', alpha=0.7)

    # Plot it!
    plt.scatter(u_rest - v_rest, v_rest - j_rest, c='blue', alpha=0.7)
    plt.xlabel('U-V (Rest Frame)')
    plt.ylabel('V-J (Rest Frame)')
    plt.xlim(0, 2.5)  # Set x-axis range from 0 to 2 (adjust as needed)
    plt.ylim(0, 2.5)
    plt.title('Rest-Frame UVJ Diagram')
    plt.grid()
    plt.show()

# Example (replace the .fits file with your own) - not real UVJ values:
rest_frame_color_color_plot('/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/glazebrook-v2_prism-clear_2565_14552.spec (3).fits', redshift=3.542, u_mag=24, v_mag=23, j_mag=21.0)

#The section in the top left is generally made up of quiescent galaxies.

def inspect_fits_file(fits_file):
    hdul = fits.open(fits_file)
    # Print the header information to see the available columns
    print(hdul[1].columns)

# Example usage:
inspect_fits_file(r"/content/drive/MyDrive/DJA-Tutorial-Example-Galaxies/glazebrook-v2_prism-clear_2565_14552.spec (3).fits")