#!/usr/bin/env python3
# encoding: utf-8
# Author: Akke Viitanen
# Email: akke.viitanen@helsinki.fi
# Date: 2023-05-04 16:17:11

"""
Test Weaver stellar mass function by visual inspection
"""


import argparse
import glob
import math
import os
import random
import re
import subprocess
import sys
import time

import astropy as ap
import astropy.coordinates as c
import astropy.units as u
import fitsio
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import scipy as sp


def double_schechter(logM, logMstar, alpha1, Phi1, alpha2, Phi2, rho):
    """Eq. 8 from Weaver+2020"""

    def fun(alpha, Phi):
        """Single Schechter function"""
        if np.isfinite(Phi) and np.isfinite(alpha):
            return Phi * (10 ** (logM - logMstar)) ** (alpha + 1)
        return np.zeros_like(logM)

    return np.log(10) * np.exp(-10 ** (logM - logMstar)) * (
        fun(alpha1, Phi1) + fun(alpha2, Phi2)
    )


def main():

    fig, axes = plt.subplots(1, 3, figsize=(3 * 6.4, 4.8), sharex=True, sharey=True)

    for ax, filename, text in zip(
        axes,
        ["total.dat", "star-forming.dat", "quiescent.dat"],
        ["Total", "Star-forming", "Quiescent"],
    ):
        ts = np.loadtxt(filename)
        n = ts.shape[0]
        logM = np.linspace(8, 12, 41)
        for i, t in enumerate(ts):
            Phi_dlogM = double_schechter(logM, *t[2:])

            norm = mpl.colors.Normalize(vmin=0, vmax=n-1)
            color = mpl.cm.viridis(norm(i))

            ax.plot(logM, Phi_dlogM, color=color)
            ax.set_xlim(8, 12.2)
            ax.set_ylim(10 ** -7, 10 ** -1.5)
            ax.semilogy()

            ax.set_xlabel("log10 M/Msun")
            ax.set_ylabel("Phi [1 / Mpc3 / dex]")

            # Do the rest only for the first drawn line
            if i > 0:
                continue
            ax.text(0.9, 0.9, text, horizontalalignment="right", transform=ax.transAxes)
    plt.show()



if __name__ == "__main__":
    main()
