#include <fstream>

#include "vtr_log.h"
#include "vtr_assert.h"
#include "vtr_math.h"

#include "globals.h"
#include "timing_util.h"
#include "timing_info.h"

double sec_to_nanosec(double seconds) {
    return 1e9 * seconds;
}

double sec_to_mhz(double seconds) {
    return (1. / seconds) / 1e6;
}

/*
 * Setup-time related
 */
tatum::TimingPathInfo find_longest_critical_path_delay(const tatum::TimingConstraints& constraints, const tatum::SetupTimingAnalyzer& setup_analyzer) {
    tatum::TimingPathInfo crit_path_info;

    auto& timing_ctx = g_vpr_ctx.timing();

    auto cpds = tatum::find_critical_paths(*timing_ctx.graph, constraints, setup_analyzer);

    //Record the maximum critical path accross all domain pairs
    for (const auto& path_info : cpds) {
        if (crit_path_info.delay() < path_info.delay() || std::isnan(crit_path_info.delay())) {
            crit_path_info = path_info;
        }
    }

    return crit_path_info;
}

tatum::TimingPathInfo find_least_slack_critical_path_delay(const tatum::TimingConstraints& constraints, const tatum::SetupTimingAnalyzer& setup_analyzer) {
    tatum::TimingPathInfo crit_path_info;

    auto& timing_ctx = g_vpr_ctx.timing();

    auto cpds = tatum::find_critical_paths(*timing_ctx.graph, constraints, setup_analyzer);

    //Record the maximum critical path accross all domain pairs
    for (const auto& path_info : cpds) {
        if (path_info.slack() < crit_path_info.slack() || std::isnan(crit_path_info.slack())) {
            crit_path_info = path_info;
        }
    }

    return crit_path_info;
}

float find_setup_total_negative_slack(const tatum::SetupTimingAnalyzer& setup_analyzer) {
    auto& timing_ctx = g_vpr_ctx.timing();

    float tns = 0.;
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : setup_analyzer.setup_slacks(node)) {
            float slack = tag.time().value();
            if (slack < 0.) {
                tns += slack;
            }
        }
    }
    return tns;
}

float find_setup_worst_negative_slack(const tatum::SetupTimingAnalyzer& setup_analyzer) {
    auto& timing_ctx = g_vpr_ctx.timing();

    float wns = 0.;
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : setup_analyzer.setup_slacks(node)) {
            float slack = tag.time().value();

            if (slack < 0.) {
                wns = std::min(wns, slack);
            }
        }
    }
    return wns;
}

float find_node_setup_slack(const tatum::SetupTimingAnalyzer& setup_analyzer, tatum::NodeId node, tatum::DomainId launch_domain, tatum::DomainId capture_domain) {
    for (const auto& tag : setup_analyzer.setup_slacks(node)) {
        if (tag.launch_clock_domain() == launch_domain && tag.capture_clock_domain() == capture_domain) {
            return tag.time().value();
        }
    }

    return NAN;
}

std::vector<HistogramBucket> create_setup_slack_histogram(const tatum::SetupTimingAnalyzer& setup_analyzer, size_t num_bins) {
    auto& timing_ctx = g_vpr_ctx.timing();

    std::vector<HistogramBucket> histogram;

    //Find the min and max slacks
    float min_slack = std::numeric_limits<float>::infinity();
    float max_slack = -std::numeric_limits<float>::infinity();
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : setup_analyzer.setup_slacks(node)) {
            float slack = tag.time().value();

            min_slack = std::min(min_slack, slack);
            max_slack = std::max(max_slack, slack);
        }
    }

    //Determine the bin size
    float range = max_slack - min_slack;
    float bin_size = range / num_bins;

    //Create the buckets
    float bucket_min = min_slack;
    for (size_t ibucket = 0; ibucket < num_bins; ++ibucket) {
        float bucket_max = bucket_min + bin_size;

        histogram.emplace_back(bucket_min, bucket_max);

        bucket_min = bucket_max;
    }

    //To avoid round-off errors we force the max value of the last bucket equal to the max slack
    histogram[histogram.size() - 1].max_value = max_slack;

    //Count the slacks into the buckets
    auto comp = [](const HistogramBucket& bucket, float slack) {
        return bucket.max_value < slack;
    };

    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : setup_analyzer.setup_slacks(node)) {
            float slack = tag.time().value();

            //Find the bucket who's max is less than the current slack

            auto iter = std::lower_bound(histogram.begin(), histogram.end(), slack, comp);
            VTR_ASSERT(iter != histogram.end());

            iter->count++;
        }
    }

    return histogram;
}

void print_setup_timing_summary(const tatum::TimingConstraints& constraints, const tatum::SetupTimingAnalyzer& setup_analyzer) {
    auto& timing_ctx = g_vpr_ctx.timing();

    auto crit_paths = tatum::find_critical_paths(*timing_ctx.graph, constraints, setup_analyzer);

    auto least_slack_cpd = find_least_slack_critical_path_delay(constraints, setup_analyzer);
    VTR_LOG("Final critical path: %g ns", sec_to_nanosec(least_slack_cpd.delay()));

    if (crit_paths.size() == 1) {
        //Fmax is only meaningful for a single-clock circuit
        VTR_LOG(", Fmax: %g MHz", sec_to_mhz(least_slack_cpd.delay()));
    }
    VTR_LOG("\n");

    VTR_LOG("Setup Worst Negative Slack (sWNS): %g ns\n", sec_to_nanosec(find_setup_worst_negative_slack(setup_analyzer)));
    VTR_LOG("Setup Total Negative Slack (sTNS): %g ns\n", sec_to_nanosec(find_setup_total_negative_slack(setup_analyzer)));
    VTR_LOG("\n");

    VTR_LOG("Setup slack histogram:\n");
    print_histogram(create_setup_slack_histogram(setup_analyzer));

    if (crit_paths.size() > 1) {
        //Multi-clock
        VTR_LOG("\n");

        //Periods per constraint
        VTR_LOG("Intra-domain critical path delays (CPDs):\n");
        for (const auto& path : crit_paths) {
            if (path.launch_domain() == path.capture_domain()) {
                VTR_LOG("  %s to %s CPD: %g ns (%g MHz)\n",
                        constraints.clock_domain_name(path.launch_domain()).c_str(),
                        constraints.clock_domain_name(path.capture_domain()).c_str(),
                        sec_to_nanosec(path.delay()),
                        sec_to_mhz(path.delay()));
            }
        }
        VTR_LOG("\n");

        VTR_LOG("Inter-domain critical path delays (CPDs):\n");
        for (const auto& path : crit_paths) {
            if (path.launch_domain() != path.capture_domain()) {
                VTR_LOG("  %s to %s CPD: %g ns (%g MHz)\n",
                        constraints.clock_domain_name(path.launch_domain()).c_str(),
                        constraints.clock_domain_name(path.capture_domain()).c_str(),
                        sec_to_nanosec(path.delay()),
                        sec_to_mhz(path.delay()));
            }
        }
        VTR_LOG("\n");

        //Slack per constraint
        VTR_LOG("Intra-domain worst setup slacks per constraint:\n");
        for (const auto& path : crit_paths) {
            if (path.launch_domain() == path.capture_domain()) {
                VTR_LOG("  %s to %s worst setup slack: %g ns\n",
                        constraints.clock_domain_name(path.launch_domain()).c_str(),
                        constraints.clock_domain_name(path.capture_domain()).c_str(),
                        sec_to_nanosec(path.slack()));
            }
        }
        VTR_LOG("\n");

        VTR_LOG("Inter-domain worst setup slacks per constraint:\n");
        for (const auto& path : crit_paths) {
            if (path.launch_domain() != path.capture_domain()) {
                VTR_LOG("  %s to %s worst setup slack: %g ns\n",
                        constraints.clock_domain_name(path.launch_domain()).c_str(),
                        constraints.clock_domain_name(path.capture_domain()).c_str(),
                        sec_to_nanosec(path.slack()));
            }
        }
    }

    //Calculate the intra-domain (i.e. same launch and capture domain) non-virtual geomean, and fanout-weighted periods
    if (crit_paths.size() > 1) {
        std::vector<double> intra_domain_cpds;
        std::vector<double> fanout_weighted_intra_domain_cpds;
        double total_intra_domain_fanout = 0.;
        auto clock_fanouts = count_clock_fanouts(*timing_ctx.graph, setup_analyzer);
        for (const auto& path : crit_paths) {
            if (path.launch_domain() == path.capture_domain() && !constraints.is_virtual_clock(path.launch_domain())) {
                intra_domain_cpds.push_back(path.delay());

                auto iter = clock_fanouts.find(path.launch_domain());
                VTR_ASSERT(iter != clock_fanouts.end());
                double fanout = iter->second;

                fanout_weighted_intra_domain_cpds.push_back(path.delay() * fanout);
                total_intra_domain_fanout += fanout;
            }
        }

        //Print multi-clock geomeans
        if (intra_domain_cpds.size() > 0) {
            VTR_LOG("\n");
            double geomean_intra_domain_cpd = vtr::geomean(intra_domain_cpds.begin(), intra_domain_cpds.end());
            VTR_LOG("Geometric mean non-virtual intra-domain period: %g ns (%g MHz)\n",
                    sec_to_nanosec(geomean_intra_domain_cpd),
                    sec_to_mhz(geomean_intra_domain_cpd));

            //Normalize weighted fanouts by total fanouts
            for (auto& weighted_cpd : fanout_weighted_intra_domain_cpds) {
                weighted_cpd /= total_intra_domain_fanout;
            }
            double fanout_weighted_geomean_intra_domain_cpd = vtr::geomean(fanout_weighted_intra_domain_cpds.begin(),
                                                                           fanout_weighted_intra_domain_cpds.end());
            VTR_LOG("Fanout-weighted geomean non-virtual intra-domain period: %g ns (%g MHz)\n",
                    sec_to_nanosec(fanout_weighted_geomean_intra_domain_cpd),
                    sec_to_mhz(fanout_weighted_geomean_intra_domain_cpd));
        }
    }
    VTR_LOG("\n");
}

/*
 * Hold-time related statistics
 */
float find_hold_total_negative_slack(const tatum::HoldTimingAnalyzer& hold_analyzer) {
    auto& timing_ctx = g_vpr_ctx.timing();

    float tns = 0.;
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : hold_analyzer.hold_slacks(node)) {
            float slack = tag.time().value();
            if (slack < 0.) {
                tns += slack;
            }
        }
    }
    return tns;
}

tatum::NodeId find_origin_node_for_hold_slack(const tatum::TimingTags::tag_range arrival_tags,
                                              const tatum::TimingTags::tag_range required_tags,
                                              float slack) {
    /*Given the slack, arrival, and required tags of a certain timing node,
     * its origin node is found*/
    for (tatum::TimingTag arrival_tag : arrival_tags) {
        for (tatum::TimingTag required_tag : required_tags) {
            if (arrival_tag.time().value() - required_tag.time().value() == slack) {
                tatum::NodeId origin_node = arrival_tag.origin_node();
                VTR_ASSERT(origin_node);
                return origin_node;
            }
        }
    }
    return tatum::NodeId::INVALID();
}

float find_total_negative_slack_within_clb_blocks(const tatum::HoldTimingAnalyzer& hold_analyzer) {
    /*Some negative slack are found within short paths in clb blocks. This cannot be optimized
     * by the router. This function is used to measure the effectiveness of the router's
     * hold time optimizer*/
    auto& timing_ctx = g_vpr_ctx.timing();
    auto& atom_ctx = g_vpr_ctx.atom();

    float slack_in_block = 0;
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : hold_analyzer.hold_slacks(node)) {
            float slack = tag.time().value();

            auto arrival_tags = hold_analyzer.hold_tags(node, tatum::TagType::DATA_ARRIVAL);
            auto required_tags = hold_analyzer.hold_tags(node, tatum::TagType::DATA_REQUIRED);
            tatum::NodeId origin_node = find_origin_node_for_hold_slack(arrival_tags, required_tags, slack);

            VTR_ASSERT(origin_node);

            /*Retrieve the source and sink clb blocks corresponding to these timing nodes*/
            AtomPinId origin_pin = atom_ctx.lookup.tnode_atom_pin(origin_node);
            AtomPinId pin = atom_ctx.lookup.tnode_atom_pin(node);
            VTR_ASSERT(origin_pin);
            VTR_ASSERT(pin);

            AtomBlockId atom_src_block = atom_ctx.nlist.pin_block(origin_pin);
            AtomBlockId atom_sink_block = atom_ctx.nlist.pin_block(pin);

            ClusterBlockId clb_src_block = atom_ctx.lookup.atom_clb(atom_src_block);
            VTR_ASSERT(clb_src_block);
            ClusterBlockId clb_sink_block = atom_ctx.lookup.atom_clb(atom_sink_block);
            VTR_ASSERT(clb_sink_block);

            const t_pb_graph_pin* sink_gpin = atom_ctx.lookup.atom_pin_pb_graph_pin(pin);
            VTR_ASSERT(sink_gpin);

            int sink_pb_route_id = sink_gpin->pin_count_in_cluster;
            ClusterNetId net_id;
            int sink_block_pin_index, sink_net_pin_index;
            std::tie(net_id, sink_block_pin_index, sink_net_pin_index) = find_pb_route_clb_input_net_pin(clb_sink_block, sink_pb_route_id);

            if (net_id == ClusterNetId::INVALID() && sink_block_pin_index == -1 && sink_net_pin_index == -1) {
                /*Does not go out of the cluster*/
                if (clb_sink_block == clb_src_block) {
                    if (slack < 0.) {
                        slack_in_block += slack;
                    }
                }
            }
        }
    }
    return slack_in_block;
}

float find_hold_worst_negative_slack(const tatum::HoldTimingAnalyzer& hold_analyzer) {
    auto& timing_ctx = g_vpr_ctx.timing();

    float wns = 0.;
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : hold_analyzer.hold_slacks(node)) {
            float slack = tag.time().value();

            if (slack < 0.) {
                wns = std::min(wns, slack);
            }
        }
    }
    return wns;
}

float find_hold_worst_slack(const tatum::HoldTimingAnalyzer& hold_analyzer, const tatum::DomainId launch, const tatum::DomainId capture) {
    auto& timing_ctx = g_vpr_ctx.timing();

    float worst_slack = std::numeric_limits<float>::infinity();
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : hold_analyzer.hold_slacks(node)) {
            if (tag.launch_clock_domain() == launch && tag.capture_clock_domain() == capture) {
                float slack = tag.time().value();

                worst_slack = std::min(worst_slack, slack);
            }
        }
    }
    return worst_slack;
}

std::vector<HistogramBucket> create_hold_slack_histogram(const tatum::HoldTimingAnalyzer& hold_analyzer, size_t num_bins) {
    auto& timing_ctx = g_vpr_ctx.timing();

    std::vector<HistogramBucket> histogram;

    //Find the min and max slacks
    float min_slack = std::numeric_limits<float>::infinity();
    float max_slack = -std::numeric_limits<float>::infinity();
    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : hold_analyzer.hold_slacks(node)) {
            float slack = tag.time().value();

            min_slack = std::min(min_slack, slack);
            max_slack = std::max(max_slack, slack);
        }
    }

    //Determine the bin size
    float range = max_slack - min_slack;
    float bin_size = range / num_bins;

    //Create the buckets
    float bucket_min = min_slack;
    for (size_t ibucket = 0; ibucket < num_bins; ++ibucket) {
        float bucket_max = bucket_min + bin_size;

        histogram.emplace_back(bucket_min, bucket_max);

        bucket_min = bucket_max;
    }

    //To avoid round-off errors we force the max value of the last bucket equal to the max slack
    histogram[histogram.size() - 1].max_value = max_slack;

    //Count the slacks into the buckets
    auto comp = [](const HistogramBucket& bucket, float slack) {
        return bucket.max_value < slack;
    };

    for (tatum::NodeId node : timing_ctx.graph->logical_outputs()) {
        for (tatum::TimingTag tag : hold_analyzer.hold_slacks(node)) {
            float slack = tag.time().value();

            //Find the bucket who's max is less than the current slack
            auto iter = std::lower_bound(histogram.begin(), histogram.end(), slack, comp);
            VTR_ASSERT(iter != histogram.end());

            //Add to bucket
            iter->count++;
        }
    }

    return histogram;
}

void print_hold_timing_summary(const tatum::TimingConstraints& constraints, const tatum::HoldTimingAnalyzer& hold_analyzer) {
    VTR_LOG("Hold Worst Negative Slack (hWNS): %g ns\n", sec_to_nanosec(find_hold_worst_negative_slack(hold_analyzer)));
    VTR_LOG("Hold Total Negative Slack (hTNS): %g ns\n", sec_to_nanosec(find_hold_total_negative_slack(hold_analyzer)));
    /*For testing*/
    //VTR_LOG("Hold Total Negative Slack within clbs: %g ns\n", sec_to_nanosec(find_total_negative_slack_within_clb_blocks(hold_analyzer)));
    VTR_LOG("\n");

    VTR_LOG("Hold slack histogram:\n");
    print_histogram(create_hold_slack_histogram(hold_analyzer));

    if (constraints.clock_domains().size() > 1) {
        //Multi-clock
        VTR_LOG("\n");

        //Slack per constraint
        VTR_LOG("Intra-domain worst hold slacks per constraint:\n");
        for (const auto& domain : constraints.clock_domains()) {
            float worst_slack = find_hold_worst_slack(hold_analyzer, domain, domain);

            if (worst_slack == std::numeric_limits<float>::infinity()) continue; //No path

            VTR_LOG("  %s to %s worst hold slack: %g ns\n",
                    constraints.clock_domain_name(domain).c_str(),
                    constraints.clock_domain_name(domain).c_str(),
                    sec_to_nanosec(worst_slack));
        }
        VTR_LOG("\n");

        VTR_LOG("Inter-domain worst hold slacks per constraint:\n");
        for (const auto& launch_domain : constraints.clock_domains()) {
            for (const auto& capture_domain : constraints.clock_domains()) {
                if (launch_domain != capture_domain) {
                    float worst_slack = find_hold_worst_slack(hold_analyzer, launch_domain, capture_domain);

                    if (worst_slack == std::numeric_limits<float>::infinity()) continue; //No path

                    VTR_LOG("  %s to %s worst hold slack: %g ns\n",
                            constraints.clock_domain_name(launch_domain).c_str(),
                            constraints.clock_domain_name(capture_domain).c_str(),
                            sec_to_nanosec(worst_slack));
                }
            }
        }
    }
    VTR_LOG("\n");
}

/*
 * General utilities
 */
std::map<tatum::DomainId, size_t> count_clock_fanouts(const tatum::TimingGraph& timing_graph, const tatum::SetupTimingAnalyzer& setup_analyzer) {
    std::map<tatum::DomainId, size_t> fanouts;
    for (tatum::NodeId node : timing_graph.nodes()) {
        tatum::NodeType type = timing_graph.node_type(node);
        if (type == tatum::NodeType::SOURCE || type == tatum::NodeType::SINK) {
            for (auto tag : setup_analyzer.setup_tags(node, tatum::TagType::DATA_ARRIVAL)) {
                fanouts[tag.launch_clock_domain()] += 1;
            }
            for (auto tag : setup_analyzer.setup_tags(node, tatum::TagType::DATA_REQUIRED)) {
                fanouts[tag.launch_clock_domain()] += 1;
            }
        }
    }

    return fanouts;
}

/*
 * Slack and criticality calculation utilities
 */

//Return the criticality of a net's pin in the CLB netlist
float calculate_clb_net_pin_criticality(const SetupTimingInfo& timing_info, const ClusteredPinAtomPinsLookup& pin_lookup, ClusterPinId clb_pin) {
    //There may be multiple atom netlist pins connected to this CLB pin
    float clb_pin_crit = 0.;
    for (const auto atom_pin : pin_lookup.connected_atom_pins(clb_pin)) {
        //Take the maximum of the atom pin criticality as the CLB pin criticality
        clb_pin_crit = std::max(clb_pin_crit, timing_info.setup_pin_criticality(atom_pin));
    }

    return clb_pin_crit;
}

//Returns the worst (maximum) criticality of the set of slack tags specified. Requires the maximum
//required time and worst slack for all domain pairs represent by the slack tags
//
// Criticality (in [0., 1.]) represents how timing-critical something is,
// 0. is non-critical and 1. is most-critical.
//
// This returns 'relaxed per constraint' criticaly as defined in:
//
//     M. Wainberg and V. Betz, "Robust Optimization of Multiple Timing Constraints,"
//         IEEE CAD, vol. 34, no. 12, pp. 1942-1953, Dec. 2015. doi: 10.1109/TCAD.2015.2440316
//
// which handles the trade-off between different timing constraints in multi-clock circuits.
//
// Note that unlike in Wainberg, we calculate the relaxed criticality as a post-processing step.

float calc_relaxed_criticality(const std::map<DomainPair, float>& domains_max_req,
                               const std::map<DomainPair, float>& domains_worst_slack,
                               const tatum::TimingTags::tag_range tags) {
    //Allowable round-off tolerance during criticality calculation
    constexpr float CRITICALITY_ROUND_OFF_TOLERANCE = 1e-4;

    //Record the maximum criticality over all the tags
    float max_crit = 0.;
    for (const auto& tag : tags) {
        VTR_ASSERT_MSG(tag.type() == tatum::TagType::SLACK, "Tags must be slacks to calculate criticality");

        float slack = tag.time().value();

        auto domain_pair = DomainPair(tag.launch_clock_domain(), tag.capture_clock_domain());

        auto iter = domains_max_req.find(domain_pair);
        VTR_ASSERT_MSG(iter != domains_max_req.end(), "Require the maximum required time for clock domain pair");
        float max_req = iter->second;

        iter = domains_worst_slack.find(domain_pair);
        VTR_ASSERT_MSG(iter != domains_worst_slack.end(), "Require the worst slack for clock domain pair");
        float worst_slack = iter->second;

        if (worst_slack < 0.) {
            //We shift slacks and required time by the most negative slack
            //**in the domain**, to ensure criticality is bounded within [0., 1.]
            //
            //This corresponds to the 'relaxed' criticality from Wainberg et. al.
            float shift = -worst_slack;
            VTR_ASSERT(shift > 0.);

            slack += shift;
            max_req += shift;
        }

        float crit = std::numeric_limits<float>::quiet_NaN();
        if (max_req > 0.) {
            //Standard case
            crit = 1. - (slack / max_req);

        } else if (max_req == 0. && slack == 0.) {
            //Special case to avoid divide by zero
            crit = 1.;
        } else {
            std::string msg = vtr::string_fmt("Invalid maximum required time %g (expected >= 0). Shifted slack was %g.", max_req, slack);
            VPR_ERROR(VPR_ERROR_TIMING, msg.c_str());
        }

        //Soft check for reasonable criticality values
        VTR_ASSERT_MSG(crit >= 0. - CRITICALITY_ROUND_OFF_TOLERANCE, "Criticality should never be negative");
        VTR_ASSERT_MSG(crit <= 1. + CRITICALITY_ROUND_OFF_TOLERANCE, "Criticality should never be greather than one");

        //Clamp criticality to [0., 1.] to correct round-off
        crit = std::max(0.f, crit);
        crit = std::min(1.f, crit);

        max_crit = std::max(max_crit, crit);
    }
    VTR_ASSERT_MSG(max_crit >= 0., "Criticality should never be negative");
    VTR_ASSERT_MSG(max_crit <= 1., "Criticality should never be greather than one");

    return max_crit;
}

void print_tatum_cpds(std::vector<tatum::TimingPathInfo> cpds) {
    for (auto path : cpds) {
        VTR_LOG("Tatum   %zu -> %zu: least_slack=%g cpd=%g\n", size_t(path.launch_domain()), size_t(path.capture_domain()), float(path.slack()), float(path.delay()));
    }
}