import numpy as np
import math
import plotly.graph_objects as go
import streamlit as st

# THIS FUNCTION CREATES SHPERES
def spheres(size, clr, dist=0): 
    # Set up 100 points. First, do angles
    theta = np.linspace(0,2*np.pi,100)
    phi = np.linspace(0,np.pi,100)
    # Set up coordinates for points on the sphere
    x0 = dist + size * np.outer(np.cos(theta),np.sin(phi))
    y0 = size * np.outer(np.sin(theta),np.sin(phi))
    z0 = size * np.outer(np.ones(100),np.cos(phi))
    # Set up trace
    trace= go.Surface(x=x0, y=y0, z=z0, colorscale=[[0,clr], [1,clr]])
    trace.update(showscale=False)
    return trace

# THIS FUNCTION CREATES ORBITS
def orbits(dist, offset=0, clr='white', wdth=2): 
    # Initialize empty lists for eac set of coordinates
    xcrd=[]
    ycrd=[]
    zcrd=[]
    # Calculate coordinates
    for i in range(0,361):
        xcrd=xcrd+[(round(np.cos(math.radians(i)),5)) * dist + offset]
        ycrd=ycrd+[(round(np.sin(math.radians(i)),5)) * dist]
        zcrd=zcrd+[0]
    trace = go.Scatter3d(x=xcrd, y=ycrd, z=zcrd, marker=dict(size=0.1), line=dict(color=clr,width=wdth))
    return trace

# THIS FUNCTION DEFINES ANNOTATIONS
def annot(xcrd, zcrd, txt, xancr='center'):
    strng=dict(showarrow=False, x=xcrd, 
               y=0, z=zcrd, text=txt, 
               xanchor=xancr, font=dict(color='#BEBEBE',size=12), opacity=0.5)
    return strng

# The Streamlit app
def main():
    st.title('Quite Accurate Solar System Visualization')
    st.write('##### Zoom in to realize how small and insignificant you are! Try to find :ringed_planet: (The planets are 1000x bigger and all orbits are further away from the sun for visual purposes.)')

    # VISUALIZE THE SOLAR SYSTEM
    diameter_km = [1391000, 4900, 12100, 12700, 6800, 143000, 120500, 51100, 49500]
    diameter = [((i / 142984) * 10) for i in diameter_km]  # Scale diameters
    distance_from_sun = [0, 264, 321, 366, 452, 1063, 1784, 3459, 5200]

    # Set up orbit traces
    traces_orbits = [
        orbits(distance_from_sun[1]), # Mercury
        orbits(distance_from_sun[2]), # Venus
        orbits(distance_from_sun[3]), # Earth
        orbits(distance_from_sun[4]), # Mars
        orbits(distance_from_sun[5]), # Jupiter
        orbits(distance_from_sun[6]), # Saturn
        orbits(distance_from_sun[7]), # Uranus
        orbits(distance_from_sun[8])  # Neptune
    ]

    # Use the same to draw a few rings for Saturn
    traces_rings_saturn = [
        orbits(15, distance_from_sun[6], '#827962', 3),
        orbits(16, distance_from_sun[6], '#827962', 3),
        orbits(17, distance_from_sun[6], '#827962', 3),
        orbits(18, distance_from_sun[6], '#827962', 3),
        orbits(19, distance_from_sun[6], '#827962', 3)
    ]

    layout=go.Layout(
        title = 'Solar System',
        showlegend=False,
        margin=dict(l=0, r=0, t=0, b=0),
        scene = dict(
            xaxis=dict(title='Distance from the Sun', 
                       titlefont_color='black', 
                       range=[-6000,6000],
                       backgroundcolor='black',
                       color='black',
                       gridcolor='black'),
            yaxis=dict(title='Distance from the Sun',
                       titlefont_color='black',
                       range=[-6000,6000],
                       backgroundcolor='black',
                       color='black',
                       gridcolor='black'
                      ),
            zaxis=dict(title='', 
                       range=[-6000,6000],
                       backgroundcolor='black',
                       color='white', 
                       gridcolor='black'
                      ),
            annotations=[
                annot(distance_from_sun[0], 200, 'Sun', xancr='left'),
                annot(distance_from_sun[1], 2, 'Mercury'),
                annot(distance_from_sun[2], 6, 'Venus'),
                annot(distance_from_sun[3], 6, 'Earth'),
                annot(distance_from_sun[4], 3, 'Mars'),
                annot(distance_from_sun[5], 70, 'Jupiter'),
                annot(distance_from_sun[6], 60, 'Saturn'),
                annot(distance_from_sun[7], 25, 'Uranus'),
                annot(distance_from_sun[8], 25, 'Neptune'),
            ]
        )
    )

    fig = go.Figure(layout=layout)
    
    # Add spheres for the Sun and planets
    fig.add_trace(spheres(diameter[0], '#ffff00', distance_from_sun[0])) # Sun
    fig.add_trace(spheres(diameter[1], '#87877d', distance_from_sun[1])) # Mercury
    fig.add_trace(spheres(diameter[2], '#d23100', distance_from_sun[2])) # Venus
    fig.add_trace(spheres(diameter[3], '#325bff', distance_from_sun[3])) # Earth
    fig.add_trace(spheres(diameter[4], '#b20000', distance_from_sun[4])) # Mars
    fig.add_trace(spheres(diameter[5], '#ebebd2', distance_from_sun[5])) # Jupiter
    fig.add_trace(spheres(diameter[6], '#ebcd82', distance_from_sun[6])) # Saturn
    fig.add_trace(spheres(diameter[7], '#37ffda', distance_from_sun[7])) # Uranus
    fig.add_trace(spheres(diameter[8], '#2500ab', distance_from_sun[8])) # Neptune

    # Add orbit traces
    for trace in traces_orbits:
        trace.line.color = '#808080'
        fig.add_trace(trace)

    # Add rings for Saturn
    for trace in traces_rings_saturn:
        fig.add_trace(trace)

    st.plotly_chart(fig)

if __name__ == "__main__":
    main()