from __future__ import division
import sys, os
import json
import numpy
from matplotlib import pyplot
from math import pi, atan2, sin, cos


def main(inpfile):
    """
    Run the thrusswork program:
    
    1. Read data from the input file
    2. Print information about the data
    3. Calculate the stress in each rod
    4. Plot the results 
    """
    # Read data from the input file
    inpdata = read_input(inpfile)
    
    # Print information about the input data
    # (this can be safely removed)
    print_info(inpdata)
    
    # Calculate stress
    stresses = calculate_stresses(inpdata)
    
    # Plot the results
    plot_results(inpdata, stresses)


def read_input(filename):
    """
    Read the input file. Since we have used an input file format that Python
    can read with built in functions we have a very easy job here
    """
    with open(filename, 'rt') as f:
        data = json.load(f)
    return data


def print_info(inpdata):
    """
    Print information about the input data
    """
    # Print info about nodes
    nodes = inpdata['nodes']
    print 'There are %d nodes in the input' % len(nodes)
    for i, node in enumerate(nodes):
        x, y = node
        print '    Node %d at x = %5.2f, y = %5.2f' % (i, x, y)
    
    # Print info about rods
    rods = inpdata['rods']
    print 'There are %d rods in the input' % len(rods)
    for i, rod in enumerate(rods):
        n1, n2 = rod
        print '    Rod %d from node %d to node %d' % (i, n1, n2)
    
    # Print info about reaction forces
    reaction_forces = inpdata['reaction_forces']
    print 'There are %d reaction forces in the input' % len(reaction_forces)
    for i, rforce in enumerate(reaction_forces):
        node, angle = rforce
        angle_deg = angle/pi*180
        print '    Force %d with angle %.2f in node %d' % (i, angle_deg, node)
    
    # Print info about loads
    loads = inpdata['loads']
    print 'There are %d loads in the input' % len(loads)
    for load in loads:
        node, angle, F = load
        angle_deg = angle/pi*180
        print '    Load with angle %.2f in node %d with magnitude %.2f N' % (angle_deg, node, F)
    
    # Print info about constants
    A = inpdata['A']
    print 'Cross section area A = %f' % A


def calculate_stresses(inpdata):
    """
    Calculate the stresses in the rods
    """
    nodes = inpdata['nodes']
    rods = inpdata['rods']
    reaction_forces = inpdata['reaction_forces']
    loads = inpdata['loads']
    area = inpdata['A']
    
    num_unknowns = len(rods) + len(reaction_forces)
    num_equations = len(nodes)*2
    print 'There are %d unknowns' % num_unknowns
    print 'There are %d equations' % num_equations
    if num_unknowns != num_equations:
        print 'ERROR, cannot solve equation system'
        exit(1)
    N = num_equations
    
    A = numpy.zeros((N, N), float)
    b = numpy.zeros(N, float)
    
    # Loop over rods and add them to the equation system
    for i, rod in enumerate(rods):
        n1, n2 = rod
        
        # Get coordinates for the nodes
        x1, y1 = nodes[n1]
        x2, y2 = nodes[n2]
        
        # Get angles 
        angle_1 = atan2(y2 - y1, x2 - x1)
        angle_2 = atan2(y1 - y2, x1 - x2)
        
        print 'Rod %d' % i
        print '    Enters node %d with angle %.2f' % (n1, angle_1/pi*180)
        print '    Enters node %d with angle %.2f' % (n2, angle_2/pi*180)
        
        # Insert into stiffness matrix for node 1
        # ...
        
        # Insert into stiffness matrix for node 2
        # ...
    
    # Loop over reaction forces and add them to the equation system
    for i, rforce in enumerate(reaction_forces):
        n1, angle_1 = rforce
        unit_force_vector1 = [cos(angle_1), sin(angle_1)]
        
        print 'Reaction force %d' % i
        print '    Enters node %d with angle %.2f' % (n1, angle_1/pi*180)
        print '    Unit force vector: [% .2f, % .2f]' % tuple(unit_force_vector1) 
        
        # Insert into stiffness matrix for node 1
        # ...
    
    # Loop over loads and add them to the equation system
    for load in loads:
        n1, angle_1, F = load
        load_vector1 = [F*cos(angle_1), F*sin(angle_1)]
        
        print 'Load %d' % i
        print '    Enters node %d with angle %.2f' % (n1, angle_1/pi*180)
        print '    Load vector: [% .2f, % .2f]' % tuple(load_vector1) 
        
        # Insert into right hand side vector
        # ...
    
    # Print equation system
    if N < 10:
        print 'A:'
        for row in A:
            for v in row:
                print '% .2f' % v,
            print
        
        print 'b:'
        for v in b:
            print '% .2f' % v,
        print
    print 'Condition number:', numpy.linalg.cond(A)
        
    # Solve the equation system
    if abs(A).max() > 0:
        forces = numpy.linalg.solve(A, b)
    else:
        print 'There is nothing in the A matrix'
        forces = numpy.zeros_like(b)
    
    # Calculate stresses
    stresses = forces / area
    
    print 'Stress range: min = %.2f, max = %.2f:' % (stresses.min(), stresses.max())
    
    return stresses


def plot_results(inpdata, stresses):
    """
    Use matplotlib to plot the results
    """
    nodes = inpdata['nodes']
    rods = inpdata['rods']
    reaction_forces = inpdata['reaction_forces']
    A = inpdata['A']
    
    # Get a color map to get plot colors for the rods
    color_map = pyplot.get_cmap('jet')
    
    max_stress = max(abs(stresses.max()), abs(stresses.min()))
    
    for i, rod in enumerate(rods):
        n1, n2 = rod
        
        # Get coordinates for the nodes
        x1, y1 = nodes[n1]
        x2, y2 = nodes[n2]
        
        # Get the stress in the rod
        stress = stresses[i]
        
        # Normalize to 0..1 range with 0 stress at 0.5
        normalized = (stress + max_stress)/(2*max_stress)
        
        # Get a color 
        color = color_map(normalized)
        
        # Plot the rod
        pyplot.plot([x1, x2], [y1, y2], color=color, linewidth=5)
        
        #print 'Rod %d has stress % .2f' % (i, stress)
    
    # Print reaction forces
    for i, rforce in enumerate(reaction_forces):
        print 'Reaction force %d is % .2f N' % (i, stresses[i+len(rods)]*A) 
        
    pyplot.xlim(-0.1, 1.1)
    pyplot.ylim(-0.1, 1.1)
    pyplot.show()
    

# Run the main() function if we are called on the command line    
if __name__ == '__main__':
    # Check if the user has given a command line argument or not
    if len(sys.argv) != 2:
        print 'ERROR, you must specify an input file'
        exit(1)
    
    # Get the name of the input file
    input_file_name = sys.argv[1]
    if not os.path.isfile(input_file_name):
        print 'ERROR, the file "%r" does not exist' % input_file_name
    
    # Run the program
    main(input_file_name)