import viktor as vkt
import itertools
import numpy as np
from pathlib import Path


class Parametrization(vkt.Parametrization):
    geometry = vkt.Tab("Geometry")
    geometry.slab = vkt.Section("Slab")
    geometry.slab.width_x = vkt.NumberField("Width in x", suffix="mm", default=6000)
    geometry.slab.width_y = vkt.NumberField("Width in y", suffix="mm", default=5000)
    geometry.slab.thickness = vkt.NumberField("Thickness", suffix="mm", default=500)

    geometry.piles = vkt.Section("Piles")
    geometry.piles.diameter = vkt.NumberField("Diameter", suffix="mm", default=500)
    geometry.piles.length = vkt.NumberField("Length", suffix="m", default=7)

    loads = vkt.Tab("Loads")
    loads.input = vkt.Section("Input")
    loads.input.uniform_load = vkt.NumberField("Uniform load", suffix="kN/m2", default=1)

    scia = vkt.Tab("SCIA")
    scia.downloads = vkt.Section("Downloads")
    scia.downloads.input_xml_btn = vkt.DownloadButton("Input .xml", method="download_scia_input_xml")
    scia.downloads.input_def_btn = vkt.DownloadButton("Input .def", method="download_scia_input_def")
    scia.downloads.input_esa_btn = vkt.DownloadButton("Input .esa", method="download_scia_input_esa")


class Foundation(vkt.Controller):
    parametrization = Parametrization

    @vkt.GeometryView("3D", x_axis_to_right=True)
    def visualize(self, params, **kwargs):
        scia_model = self.create_scia_model(params)
        geometries = self.create_visualization_geometries(params, scia_model)
        return vkt.GeometryResult(geometries)

    @vkt.GeometryAndDataView("SCIA result", duration_guess=60, x_axis_to_right=True)
    def run_scia(self, params, **kwargs):
        scia_model = self.create_scia_model(params)

        # create input files
        input_xml, input_def = scia_model.generate_xml_input()
        input_esa = self.get_scia_input_esa()

        # analyze SCIA model
        scia_analysis = vkt.scia.SciaAnalysis(input_xml, input_def, input_esa)
        scia_analysis.execute(300)  # timeout after 5 minutes
        scia_result = scia_analysis.get_xml_output_file()

        # parse analysis result
        result = vkt.scia.OutputFileParser.get_result(scia_result, "Reactions", parent='Combinations - C1')
        reactions = result['Nodal reactions']

        max_rz = float(max(reactions['R_z']))
        data_result = vkt.DataGroup(
            vkt.DataItem('SCIA results', ' ', subgroup=vkt.DataGroup(
                vkt.DataItem('Maximum pile reaction', max_rz, suffix='N', number_of_decimals=2)
            ))
        )

        geometries = self.create_visualization_geometries(params, scia_model)
        return vkt.GeometryAndDataResult(geometries, data_result)

    def download_scia_input_esa(self, params, **kwargs):
        scia_input_esa = self.get_scia_input_esa()
        return vkt.DownloadResult(scia_input_esa, 'model.esa')

    def download_scia_input_xml(self, params, **kwargs):
        scia_model = self.create_scia_model(params)
        input_xml, _ = scia_model.generate_xml_input()
        return vkt.DownloadResult(input_xml, 'test.xml')

    def download_scia_input_def(self, params, **kwargs):
        m = vkt.scia.SciaModel()
        _, input_def = m.generate_xml_input()
        return vkt.DownloadResult(input_def, 'viktor.xml.def')

    def get_scia_input_esa(self) -> vkt.File:
        return vkt.File.from_path(Path(__file__).parent / 'model.esa')

    def create_scia_model(self, params) -> vkt.scia.Model:
        model = vkt.scia.Model()

        '''
        STRUCTURE
        '''
        # create nodes at the slab corners
        width_x = params.geometry.slab.width_x * 1e-03
        width_y = params.geometry.slab.width_y * 1e-03
        n1 = model.create_node('n1', 0, 0, 0)  # origin
        n2 = model.create_node('n2', 0, width_y, 0)
        n3 = model.create_node('n3', width_x, width_y, 0)
        n4 = model.create_node('n4', width_x, 0, 0)

        # create the pile nodes
        number_of_piles_x = 4
        number_of_piles_y = 3
        pile_edge_distance = 0.3
        pile_length = params.geometry.piles.length

        start_x = pile_edge_distance
        end_x = width_x - pile_edge_distance
        x = np.linspace(start_x, end_x, number_of_piles_x)
        start_y = pile_edge_distance
        end_y = width_y - pile_edge_distance
        y = np.linspace(start_y, end_y, number_of_piles_y)
        pile_positions = np.array(list(itertools.product(x, y)))

        pile_top_nodes = []
        pile_bottom_nodes = []
        for pile_id, (pile_x, pile_y) in enumerate(pile_positions, 1):
            n_top = model.create_node(f'K:p{pile_id}_t', pile_x, pile_y, 0)
            n_bottom = model.create_node(f'K:p{pile_id}_b', pile_x, pile_y, -pile_length)
            pile_top_nodes.append(n_top)
            pile_bottom_nodes.append(n_bottom)

        # create pile beams
        pile_diameter = params.geometry.piles.diameter * 1e-03
        material = vkt.scia.Material(0, 'C30/37')
        cross_section = model.create_circular_cross_section('concrete_pile', material, pile_diameter)
        pile_beams = []
        for pile_id, (n_top, n_bottom) in enumerate(zip(pile_top_nodes, pile_bottom_nodes), 1):
            pile_beam = model.create_beam(n_top, n_bottom, cross_section)
            pile_beams.append(pile_beam)

        # create the concrete slab
        material = vkt.scia.Material(0, 'concrete_slab')
        thickness = params.geometry.slab.thickness * 1e-03
        corner_nodes = [n1, n2, n3, n4]
        slab = model.create_plane(corner_nodes, thickness, name='foundation slab', material=material)

        '''
        SUPPORTS
        '''

        # create pile point supports
        freedom_v = (
            vkt.scia.PointSupport.Freedom.FREE,
            vkt.scia.PointSupport.Freedom.FREE,
            vkt.scia.PointSupport.Freedom.FLEXIBLE,
            vkt.scia.PointSupport.Freedom.FREE,
            vkt.scia.PointSupport.Freedom.FREE,
            vkt.scia.PointSupport.Freedom.FREE
        )
        kv = 400 * 1e06
        stiffness_v = (0, 0, kv, 0, 0, 0)
        for pile_id, pile_beam in enumerate(pile_beams, 1):
            n_bottom = pile_beam.end_node
            model.create_point_support(f'Sn:p{pile_id}', n_bottom, vkt.scia.PointSupport.Type.STANDARD,
                                       freedom_v, stiffness_v, vkt.scia.PointSupport.CSys.GLOBAL)

        # create pile line supports
        kh = 10 * 1e06
        for pile_id, pile_beam in enumerate(pile_beams, 1):
            model.create_line_support_on_beam(pile_beam, x=vkt.scia.LineSupport.Freedom.FLEXIBLE, stiffness_x=kh,
                                              y=vkt.scia.LineSupport.Freedom.FLEXIBLE, stiffness_y=kh,
                                              z=vkt.scia.LineSupport.Freedom.FREE, rx=vkt.scia.LineSupport.Freedom.FREE,
                                              ry=vkt.scia.LineSupport.Freedom.FREE, rz=vkt.scia.LineSupport.Freedom.FREE,
                                              c_sys=vkt.scia.LineSupport.CSys.GLOBAL)

        # create the slab supports
        stiffness_x = 50 * 1e06
        stiffness_y = 50 * 1e06
        for edge in (1, 3):
            model.create_line_support_on_plane((slab, edge),
                                               x=vkt.scia.LineSupport.Freedom.FLEXIBLE, stiffness_x=stiffness_x,
                                               y=vkt.scia.LineSupport.Freedom.FREE,
                                               z=vkt.scia.LineSupport.Freedom.FREE,
                                               rx=vkt.scia.LineSupport.Freedom.FREE,
                                               ry=vkt.scia.LineSupport.Freedom.FREE,
                                               rz=vkt.scia.LineSupport.Freedom.FREE)
        for edge in (2, 4):
            model.create_line_support_on_plane((slab, edge),
                                               x=vkt.scia.LineSupport.Freedom.FREE,
                                               y=vkt.scia.LineSupport.Freedom.FLEXIBLE, stiffness_y=stiffness_y,
                                               z=vkt.scia.LineSupport.Freedom.FREE,
                                               rx=vkt.scia.LineSupport.Freedom.FREE,
                                               ry=vkt.scia.LineSupport.Freedom.FREE,
                                               rz=vkt.scia.LineSupport.Freedom.FREE)

        '''
        SETS
        '''
        # create the load group
        lg = model.create_load_group('LG1', vkt.scia.LoadGroup.LoadOption.VARIABLE, vkt.scia.LoadGroup.RelationOption.STANDARD,
                                     vkt.scia.LoadGroup.LoadTypeOption.CAT_G)

        # create the load case
        lc = model.create_variable_load_case('LC1', 'first load case', lg, vkt.scia.LoadCase.VariableLoadType.STATIC,
                                             vkt.scia.LoadCase.Specification.STANDARD, vkt.scia.LoadCase.Duration.SHORT)

        # create the load combination
        load_cases = {
            lc: 1
        }

        model.create_load_combination('C1', vkt.scia.LoadCombination.Type.ENVELOPE_SERVICEABILITY, load_cases)

        '''
        LOADS
        '''
        # create the load
        force = params.loads.input.uniform_load * 1e03
        force *= -1  # in negative Z-direction
        model.create_surface_load('SF:1', lc, slab, vkt.scia.SurfaceLoad.Direction.Z, vkt.scia.SurfaceLoad.Type.FORCE, force,
                                  vkt.scia.SurfaceLoad.CSys.GLOBAL, vkt.scia.SurfaceLoad.Location.LENGTH)

        return model

    def create_visualization_geometries(self, params, scia_model):
        geometries = []
        for node in scia_model.nodes:
            node_obj = vkt.Sphere(vkt.Point(node.x, node.y, node.z), params.geometry.slab.width_y * 1e-05)
            node_obj.material = vkt.Material('node', color=vkt.Color(0, 255, 0))
            geometries.append(node_obj)

        # pile beams
        pile_diameter = params.geometry.piles.diameter * 1e-03
        for beam in scia_model.beams:
            point_top = vkt.Point(beam.begin_node.x, beam.begin_node.y, beam.begin_node.z)
            point_bottom = vkt.Point(beam.end_node.x, beam.end_node.y, beam.end_node.z)
            beam_obj = vkt.CircularExtrusion(pile_diameter, vkt.Line(point_top, point_bottom))
            beam_obj.material = vkt.Material('beam', roughness=1, opacity=0.3)
            geometries.append(beam_obj)

        corner_points = [
            vkt.Point(scia_model.nodes[0].x, scia_model.nodes[0].y, scia_model.nodes[0].z),
            vkt.Point(scia_model.nodes[1].x, scia_model.nodes[1].y, scia_model.nodes[1].z),
            vkt.Point(scia_model.nodes[2].x, scia_model.nodes[2].y, scia_model.nodes[2].z),
            vkt.Point(scia_model.nodes[3].x, scia_model.nodes[3].y, scia_model.nodes[3].z),
            vkt.Point(scia_model.nodes[0].x, scia_model.nodes[0].y, scia_model.nodes[0].z)
        ]

        thickness = params.geometry.slab.thickness * 1e-03
        slab_obj = vkt.Extrusion(corner_points, vkt.Line(vkt.Point(0, 0, -thickness / 2), vkt.Point(0, 0, thickness / 2)))
        slab_obj.material = vkt.Material('slab', roughness=1, opacity=0.3)
        geometries.append(slab_obj)

        return geometries


initial_entities = [
    vkt.InitialEntity("Foundation", "Example")
]
